Merge pull request #516 from databrickslabs/buffer-params

Buffer params

R/sparklyr-mosaic/sparklyrMosaic/tests/testthat/testVectorFunctions.R

@@ -15,6 +15,7 @@ test_that("scalar vector functions behave as intended", {
     st_length = st_length(wkt),
     st_perimeter = st_perimeter(wkt),
     st_buffer = st_buffer(wkt, as.double(1.1)),
+    st_buffer_optparams = st_buffer(wkt, as.double(1.1), "endcap=square quad_segs=2"),
     st_bufferloop = st_bufferloop(wkt, as.double(1.1), as.double(1.2)),
     st_convexhull = st_convexhull(wkt),
     st_dump = st_dump(wkt),
@@ -122,4 +123,4 @@ test_that("aggregate vector functions behave as intended", {
   expect_true(sdf.intersection %>% head(1) %>% sdf_collect %>% .$comparison_intersection)
 
 
-})
+})

docs/source/api/spatial-functions.rst

@@ -125,11 +125,17 @@ st_buffer
 .. function:: st_buffer(col, radius)
 
     Buffer the input geometry by radius :code:`radius` and return a new, buffered geometry.
+    The optional parameter buffer_style_parameters='quad_segs=# endcap=round|flat|square' where "#"
+    is the number of line segments used to approximate a quarter circle (default is 8); and endcap
+    style for line features is one of listed (default="round")
+
 
     :param col: Geometry
     :type col: Column
     :param radius: Double
     :type radius: Column (DoubleType)
+    :param buffer_style_parameters: String
+    :type buffer_style_parameters: Column (StringType)
     :rtype: Column: Geometry
 
     :example:
@@ -1535,7 +1541,7 @@ st_transform
 
     df <- createDataFrame(data.frame(wkt = "MULTIPOINT ((10 40), (40 30), (20 20), (30 10))"))
     df <- withColumn(df, 'geom', st_setsrid(st_asgeojson(column('wkt')), lit(4326L)))
-    >>>
+
     showDF(select(df, st_astext(st_transform(column('geom'), lit(3857L)))), truncate=F)
     +--------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
     |convert_to(st_transform(geom, 3857))                                                                                                                                      |

python/mosaic/api/functions.py

@@ -197,7 +197,11 @@ def st_concavehull(
     )
 
 
-def st_buffer(geom: ColumnOrName, radius: ColumnOrName) -> Column:
+def st_buffer(
+        geom: ColumnOrName,
+        radius: ColumnOrName,
+        buffer_style_parameters: Any = "",
+) -> Column:
     """
     Compute the buffered geometry based on geom and radius.
 
@@ -207,15 +211,27 @@ def st_buffer(geom: ColumnOrName, radius: ColumnOrName) -> Column:
         The input geometry
     radius : Column
         The radius of buffering
+    buffer_style_parameters : Column
+        "quad_segs=# endcap=round|flat|square" where "#" is the number of line segments used to
+        approximate a quarter circle (default is 8); and endcap style for line features is one of
+        listed (default="round")
+
 
     Returns
     -------
     Column
         A geometry
 
     """
+
+    if isinstance(buffer_style_parameters, str):
+        buffer_style_parameters = lit(buffer_style_parameters)
+
     return config.mosaic_context.invoke_function(
-        "st_buffer", pyspark_to_java_column(geom), pyspark_to_java_column(radius)
+        "st_buffer",
+        pyspark_to_java_column(geom),
+        pyspark_to_java_column(radius),
+        pyspark_to_java_column(buffer_style_parameters),
     )
 
 

python/test/test_vector_functions.py

@@ -59,6 +59,9 @@ def test_st_bindings_happy_flow(self):
             df.withColumn("st_area", api.st_area("wkt"))
             .withColumn("st_length", api.st_length("wkt"))
             .withColumn("st_buffer", api.st_buffer("wkt", lit(1.1)))
+            .withColumn(
+                "st_buffer_optparams",
+                api.st_buffer("wkt", lit(1.1), lit("endcap=square quad_segs=2")))
             .withColumn("st_bufferloop", api.st_bufferloop("wkt", lit(1.1), lit(1.2)))
             .withColumn("st_perimeter", api.st_perimeter("wkt"))
             .withColumn("st_convexhull", api.st_convexhull("wkt"))

src/main/scala/com/databricks/labs/mosaic/core/geometry/MosaicGeometry.scala

@@ -59,6 +59,8 @@ trait MosaicGeometry extends GeometryWriter with Serializable {
 
     def buffer(distance: Double): MosaicGeometry
 
+    def buffer(distance: Double, bufferStyleParameters: String = ""): MosaicGeometry
+
     def bufferCapStyle(distance: Double, capStyle: String): MosaicGeometry
 
     def simplify(tolerance: Double): MosaicGeometry

src/main/scala/com/databricks/labs/mosaic/core/geometry/MosaicGeometryJTS.scala

@@ -73,7 +73,36 @@ abstract class MosaicGeometryJTS(geom: Geometry) extends MosaicGeometry {
     override def envelope: MosaicGeometryJTS = MosaicGeometryJTS(geom.getEnvelope)
 
     override def buffer(distance: Double): MosaicGeometryJTS = {
-        val buffered = geom.buffer(distance)
+        buffer(distance, "")
+    }
+
+    override def buffer(distance: Double, bufferStyleParameters: String): MosaicGeometryJTS = {
+
+        val gBuf = new BufferOp(geom)
+
+        if (bufferStyleParameters contains "=") {
+            val params = bufferStyleParameters
+                .split(" ")
+                .map(_.split("="))
+                .map { case Array(k, v) => (k, v) }
+                .toMap
+
+            if (params.contains("endcap")) {
+                val capStyle = params.getOrElse("endcap", "")
+                val capStyleConst = capStyle match {
+                    case "round"  => BufferParameters.CAP_ROUND
+                    case "flat"   => BufferParameters.CAP_FLAT
+                    case "square" => BufferParameters.CAP_SQUARE
+                    case _        => BufferParameters.CAP_ROUND
+                }
+                gBuf.setEndCapStyle(capStyleConst)
+            }
+            if (params.contains("quad_segs")) {
+                val quadSegs = params.getOrElse("quad_segs", "8")
+                gBuf.setQuadrantSegments(quadSegs.toInt)
+            }
+        }
+        val buffered = gBuf.getResultGeometry(distance)
         buffered.setSRID(geom.getSRID)
         MosaicGeometryJTS(buffered)
     }

src/main/scala/com/databricks/labs/mosaic/expressions/geometry/ST_Buffer.scala

@@ -2,40 +2,48 @@ package com.databricks.labs.mosaic.expressions.geometry
 
 import com.databricks.labs.mosaic.core.geometry.MosaicGeometry
 import com.databricks.labs.mosaic.expressions.base.WithExpressionInfo
-import com.databricks.labs.mosaic.expressions.geometry.base.UnaryVector1ArgExpression
+import com.databricks.labs.mosaic.expressions.geometry.base.UnaryVector2ArgExpression
 import com.databricks.labs.mosaic.functions.MosaicExpressionConfig
 import org.apache.spark.sql.adapters.Column
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.FunctionBuilder
 import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext
 import org.apache.spark.sql.types.DataType
+import org.apache.spark.unsafe.types.UTF8String
+import org.apache.spark.sql.functions._
 
 /**
   * SQL expression that returns the input geometry buffered by the radius.
   * @param inputGeom
   *   Expression containing the geometry.
   * @param radiusExpr
   *   The radius of the buffer.
+  * @param bufferStyleParametersExpr
+  *   'quad_segs=# endcap=round|flat|square' where "#" is the number of line
+  *   segments used to approximate a quarter circle (default is 8); and endcap
+  *   style for line features is one of listed (default="round")
   * @param expressionConfig
   *   Mosaic execution context, e.g. geometryAPI, indexSystem, etc. Additional
   *   arguments for the expression (expressionConfigs).
   */
 case class ST_Buffer(
     inputGeom: Expression,
     radiusExpr: Expression,
+    bufferStyleParametersExpr: Expression = lit("").expr,
     expressionConfig: MosaicExpressionConfig
-) extends UnaryVector1ArgExpression[ST_Buffer](inputGeom, radiusExpr, returnsGeometry = true, expressionConfig) {
+) extends UnaryVector2ArgExpression[ST_Buffer](inputGeom, radiusExpr, bufferStyleParametersExpr, returnsGeometry = true, expressionConfig) {
 
     override def dataType: DataType = inputGeom.dataType
 
-    override def geometryTransform(geometry: MosaicGeometry, arg: Any): Any = {
-        val radius = arg.asInstanceOf[Double]
-        geometry.buffer(radius)
+    override def geometryTransform(geometry: MosaicGeometry, arg1: Any, arg2: Any): Any = {
+        val radius = arg1.asInstanceOf[Double]
+        val bufferStyleParameters = arg2.asInstanceOf[UTF8String].toString
+        geometry.buffer(radius, bufferStyleParameters)
     }
 
-    override def geometryCodeGen(geometryRef: String, argRef: String, ctx: CodegenContext): (String, String) = {
+    override def geometryCodeGen(geometryRef: String, argRef1: String, argRef2: String, ctx: CodegenContext): (String, String) = {
         val resultRef = ctx.freshName("result")
-        val code = s"""$mosaicGeomClass $resultRef = $geometryRef.buffer($argRef);"""
+        val code = s"""$mosaicGeomClass $resultRef = $geometryRef.buffer($argRef1, $argRef2.toString());"""
         (code, resultRef)
     }
 
@@ -56,7 +64,11 @@ object ST_Buffer extends WithExpressionInfo {
           |  """.stripMargin
 
     override def builder(expressionConfig: MosaicExpressionConfig): FunctionBuilder = { (children: Seq[Expression]) =>
-        ST_Buffer(children.head, Column(children(1)).cast("double").expr, expressionConfig)
+        if (children.size == 2) {
+            ST_Buffer(children.head, Column(children(1)).cast("double").expr, lit("").expr, expressionConfig)
+        } else if (children.size == 3) {
+            ST_Buffer(children.head, Column(children(1)).cast("double").expr, Column(children(2)).cast("string").expr, expressionConfig)
+        } else throw new Exception("unexpected number of arguments")
     }
 
 }

src/main/scala/com/databricks/labs/mosaic/functions/MosaicContext.scala

@@ -561,10 +561,14 @@ class MosaicContext(indexSystem: IndexSystem, geometryAPI: GeometryAPI) extends
         /** Spatial functions */
         def flatten_polygons(geom: Column): Column = ColumnAdapter(FlattenPolygons(geom.expr, geometryAPI.name))
         def st_area(geom: Column): Column = ColumnAdapter(ST_Area(geom.expr, expressionConfig))
-        def st_buffer(geom: Column, radius: Column): Column =
-            ColumnAdapter(ST_Buffer(geom.expr, radius.cast("double").expr, expressionConfig))
-        def st_buffer(geom: Column, radius: Double): Column =
-            ColumnAdapter(ST_Buffer(geom.expr, lit(radius).cast("double").expr, expressionConfig))
+        def st_buffer(geom: Column, radius: Column): Column = st_buffer(geom, radius, lit(""))
+        def st_buffer(geom: Column, radius: Double): Column = st_buffer(geom, lit(radius), lit(""))
+        def st_buffer(geom: Column, radius: Column, buffer_style_parameters: Column): Column =
+            ColumnAdapter(ST_Buffer(geom.expr, radius.cast("double").expr, buffer_style_parameters.cast("string").expr, expressionConfig))
+        def st_buffer(geom: Column, radius: Double, buffer_style_parameters: Column): Column =
+            ColumnAdapter(
+              ST_Buffer(geom.expr, lit(radius).cast("double").expr, lit(buffer_style_parameters).cast("string").expr, expressionConfig)
+            )
         def st_bufferloop(geom: Column, r1: Column, r2: Column): Column =
             ColumnAdapter(ST_BufferLoop(geom.expr, r1.cast("double").expr, r2.cast("double").expr, expressionConfig))
         def st_bufferloop(geom: Column, r1: Double, r2: Double): Column =