Skip to content

Commit

Permalink
Merge pull request #494 from databrickslabs/feature/st_concavehull_is…
Browse files Browse the repository at this point in the history
…sue117

Add support for ST_ConcaveHull.
  • Loading branch information
Milos Colic authored Jan 12, 2024
2 parents cee449c + 4d91d00 commit 8fa5d6a
Show file tree
Hide file tree
Showing 9 changed files with 280 additions and 14 deletions.
68 changes: 54 additions & 14 deletions python/mosaic/api/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
"st_length",
"st_perimeter",
"st_convexhull",
"st_concavehull",
"st_buffer",
"st_bufferloop",
"st_dimension",
Expand Down Expand Up @@ -154,6 +155,45 @@ def st_convexhull(geom: ColumnOrName) -> Column:
)


def st_concavehull(geom: ColumnOrName, concavity: ColumnOrName, has_holes: Any = False) -> Column:
"""
Compute the concave hull of a geometry or multi-geometry object.
It uses lengthRatio and
allowHoles to determine the concave hull. lengthRatio is the ratio of the
length of the concave hull to the length of the convex hull. If set to 1,
this is the same as the convex hull. If set to 0, this is the same as the
bounding box. AllowHoles is a boolean that determines whether the concave
hull can have holes. If set to true, the concave hull can have holes. If set
to false, the concave hull will not have holes. (For PostGIS, the default is
false.)
Parameters
----------
geom : Column
The input geometry
concavity : Column
The concavity of the hull
has_holes : Column
Whether the hull has holes
Returns
-------
Column
A polygon
"""

if type(has_holes) == bool:
has_holes = lit(has_holes)

return config.mosaic_context.invoke_function(
"st_concavehull",
pyspark_to_java_column(geom),
pyspark_to_java_column(concavity),
pyspark_to_java_column(has_holes)
)


def st_buffer(geom: ColumnOrName, radius: ColumnOrName) -> Column:
"""
Compute the buffered geometry based on geom and radius.
Expand All @@ -177,7 +217,7 @@ def st_buffer(geom: ColumnOrName, radius: ColumnOrName) -> Column:


def st_bufferloop(
geom: ColumnOrName, inner_radius: ColumnOrName, outer_radius: ColumnOrName
geom: ColumnOrName, inner_radius: ColumnOrName, outer_radius: ColumnOrName
) -> Column:
"""
Compute the buffered geometry loop (hollow ring) based on geom and provided radius-es.
Expand Down Expand Up @@ -323,7 +363,7 @@ def st_transform(geom: ColumnOrName, srid: ColumnOrName) -> Column:


def st_hasvalidcoordinates(
geom: ColumnOrName, crs: ColumnOrName, which: ColumnOrName
geom: ColumnOrName, crs: ColumnOrName, which: ColumnOrName
) -> Column:
"""
Checks if all points in geometry are valid with respect to crs bounds.
Expand Down Expand Up @@ -530,7 +570,7 @@ def st_distance(geom1: ColumnOrName, geom2: ColumnOrName) -> Column:


def st_haversine(
lat1: ColumnOrName, lng1: ColumnOrName, lat2: ColumnOrName, lng2: ColumnOrName
lat1: ColumnOrName, lng1: ColumnOrName, lat2: ColumnOrName, lng2: ColumnOrName
) -> Column:
"""
Compute the haversine distance in kilometers between two latitude/longitude pairs.
Expand Down Expand Up @@ -682,7 +722,7 @@ def st_unaryunion(geom: ColumnOrName) -> Column:


def st_updatesrid(
geom: ColumnOrName, srcSRID: ColumnOrName, destSRID: ColumnOrName
geom: ColumnOrName, srcSRID: ColumnOrName, destSRID: ColumnOrName
) -> Column:
"""
Updates the SRID of the input geometry `geom` from `srcSRID` to `destSRID`.
Expand Down Expand Up @@ -951,7 +991,7 @@ def grid_boundary(index_id: ColumnOrName, format_name: ColumnOrName) -> Column:


def grid_longlatascellid(
lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
) -> Column:
"""
Returns the grid's cell ID associated with the input `lng` and `lat` coordinates at a given grid `resolution`.
Expand Down Expand Up @@ -1019,7 +1059,7 @@ def grid_polyfill(geom: ColumnOrName, resolution: ColumnOrName) -> Column:


def grid_tessellate(
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
) -> Column:
"""
Generates:
Expand Down Expand Up @@ -1054,7 +1094,7 @@ def grid_tessellate(


def grid_tessellateexplode(
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
) -> Column:
"""
Generates:
Expand Down Expand Up @@ -1214,7 +1254,7 @@ def grid_cellkloopexplode(cellid: ColumnOrName, k: ColumnOrName) -> Column:


def grid_geometrykring(
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
) -> Column:
"""
Returns the k-ring of cells around the input geometry.
Expand All @@ -1239,7 +1279,7 @@ def grid_geometrykring(


def grid_geometrykloop(
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
) -> Column:
"""
Returns the k loop (hollow ring) of cells around the input geometry.
Expand All @@ -1264,7 +1304,7 @@ def grid_geometrykloop(


def grid_geometrykringexplode(
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
) -> Column:
"""
Returns the exploded k-ring of cells around the input geometry.
Expand All @@ -1289,7 +1329,7 @@ def grid_geometrykringexplode(


def grid_geometrykloopexplode(
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
geom: ColumnOrName, resolution: ColumnOrName, k: ColumnOrName
) -> Column:
"""
Returns the exploded k loop (hollow ring) of cells around the input geometry.
Expand Down Expand Up @@ -1336,7 +1376,7 @@ def point_index_geom(geom: ColumnOrName, resolution: ColumnOrName) -> Column:


def point_index_lonlat(
lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
lon: ColumnOrName, lat: ColumnOrName, resolution: ColumnOrName
) -> Column:
"""
[Deprecated] alias for `grid_longlatascellid`
Expand Down Expand Up @@ -1393,7 +1433,7 @@ def polyfill(geom: ColumnOrName, resolution: ColumnOrName) -> Column:


def mosaic_explode(
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
) -> Column:
"""
[Deprecated] alias for `grid_tessellateexplode`
Expand Down Expand Up @@ -1428,7 +1468,7 @@ def mosaic_explode(


def mosaicfill(
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
geom: ColumnOrName, resolution: ColumnOrName, keep_core_geometries: Any = True
) -> Column:
"""
[Deprecated] alias for `grid_tessellate`
Expand Down
1 change: 1 addition & 0 deletions python/test/test_vector_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ def test_st_bindings_happy_flow(self):
.withColumn("st_buffer", api.st_bufferloop("wkt", lit(1.1), lit(1.2)))
.withColumn("st_perimeter", api.st_perimeter("wkt"))
.withColumn("st_convexhull", api.st_convexhull("wkt"))
.withColumn("st_concavehull", api.st_concavehull("wkt", lit(0.5)))
.withColumn("st_dump", api.st_dump("wkt"))
.withColumn("st_translate", api.st_translate("wkt", lit(1), lit(1)))
.withColumn("st_scale", api.st_scale("wkt", lit(1), lit(1)))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,9 @@ trait MosaicGeometry extends GeometryWriter with Serializable {

def convexHull: MosaicGeometry

// Allow holes is set to false by default to match the behavior of the POSTGIS implementation
def concaveHull(lengthRatio: Double, allow_holes: Boolean = false): MosaicGeometry

def minMaxCoord(dimension: String, func: String): Double = {
val coordArray = this.getShellPoints.map(shell => {
val unitArray = dimension.toUpperCase(Locale.ROOT) match {
Expand Down
Empty file.
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import com.databricks.labs.mosaic.core.types.model.GeometryTypeEnum
import com.databricks.labs.mosaic.core.types.model.GeometryTypeEnum._
import com.esotericsoftware.kryo.Kryo
import org.apache.spark.sql.catalyst.InternalRow
import org.locationtech.jts.algorithm.hull.ConcaveHull
import org.locationtech.jts.geom.{Geometry, GeometryCollection, GeometryFactory}
import org.locationtech.jts.geom.util.AffineTransformation
import org.locationtech.jts.io._
Expand Down Expand Up @@ -182,6 +183,12 @@ abstract class MosaicGeometryJTS(geom: Geometry) extends MosaicGeometry {
MosaicGeometryJTS(convexHull)
}

override def concaveHull(lengthRatio: Double, allow_holes: Boolean = false): MosaicGeometryJTS = {
val concaveHull = ConcaveHull.concaveHullByLengthRatio(geom, lengthRatio, allow_holes)
concaveHull.setSRID(geom.getSRID)
MosaicGeometryJTS(concaveHull)
}

override def unaryUnion: MosaicGeometryJTS = {
val unaryUnion = geom.union()
unaryUnion.setSRID(geom.getSRID)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
package com.databricks.labs.mosaic.expressions.geometry

import com.databricks.labs.mosaic.core.geometry.MosaicGeometry
import com.databricks.labs.mosaic.expressions.base.{GenericExpressionFactory, WithExpressionInfo}
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

/**
* Returns the concave hull for a given geometry. It uses lengthRatio and
* allowHoles to determine the concave hull. lengthRatio is the ratio of the
* length of the concave hull to the length of the convex hull. If set to 1,
* this is the same as the convex hull. If set to 0, this is the same as the
* bounding box. AllowHoles is a boolean that determines whether the concave
* hull can have holes. If set to true, the concave hull can have holes. If set
* to false, the concave hull will not have holes. (For PostGIS, the default is
* false.)
* @param inputGeom
* The input geometry.
* @param expressionConfig
* Additional arguments for the expression (expressionConfigs).
*/
case class ST_ConcaveHull(
inputGeom: Expression,
lengthRatio: Expression,
allowHoles: Expression,
expressionConfig: MosaicExpressionConfig
) extends UnaryVector2ArgExpression[ST_ConcaveHull](
inputGeom,
lengthRatio,
allowHoles,
returnsGeometry = true,
expressionConfig
) {

override def dataType: DataType = inputGeom.dataType

override def geometryTransform(geometry: MosaicGeometry, arg1: Any, arg2: Any): Any = {
val lenRatio = arg1.asInstanceOf[Double]
val allowHoles = arg2.asInstanceOf[Boolean]
geometry.concaveHull(lenRatio, allowHoles)
}

override def geometryCodeGen(geometryRef: String, arg1Ref: String, arg2Ref: String, ctx: CodegenContext): (String, String) = {
val convexHull = ctx.freshName("concaveHull")
val code = s"""$mosaicGeomClass $convexHull = $geometryRef.concaveHull($arg1Ref, $arg2Ref);"""
(code, convexHull)
}

}

/** Expression info required for the expression registration for spark SQL. */
object ST_ConcaveHull extends WithExpressionInfo {

override def name: String = "st_concavehull"

override def usage: String = "_FUNC_(expr1, expr2, expr3) - Returns the concave hull for a given geometry with or without holes."

override def example: String =
"""
| Examples:
| > SELECT _FUNC_(a, 0.1, false);
| {"POLYGON (( 0 0, 1 0, 1 1, 0 1 ))"}
| """.stripMargin

override def builder(expressionConfig: MosaicExpressionConfig): FunctionBuilder = { (children: Seq[Expression]) =>
GenericExpressionFactory.construct[ST_ConcaveHull](
Array(children.head, Column(children(1)).cast("double").expr, children(2)),
expressionConfig
)
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,7 @@ class MosaicContext(indexSystem: IndexSystem, geometryAPI: GeometryAPI) extends
mosaicRegistry.registerExpression[ST_Centroid](expressionConfig)
mosaicRegistry.registerExpression[ST_Contains](expressionConfig)
mosaicRegistry.registerExpression[ST_ConvexHull](expressionConfig)
mosaicRegistry.registerExpression[ST_ConcaveHull](expressionConfig)
mosaicRegistry.registerExpression[ST_Distance](expressionConfig)
mosaicRegistry.registerExpression[ST_Difference](expressionConfig)
mosaicRegistry.registerExpression[ST_Dimension](expressionConfig)
Expand Down Expand Up @@ -560,6 +561,10 @@ class MosaicContext(indexSystem: IndexSystem, geometryAPI: GeometryAPI) extends
ColumnAdapter(ST_BufferCapStyle(geom.expr, lit(radius).cast("double").expr, lit(capStyle).expr, expressionConfig))
def st_centroid(geom: Column): Column = ColumnAdapter(ST_Centroid(geom.expr, expressionConfig))
def st_convexhull(geom: Column): Column = ColumnAdapter(ST_ConvexHull(geom.expr, expressionConfig))
def st_concavehull(geom: Column, concavity: Column, allowHoles: Column): Column =
ColumnAdapter(ST_ConcaveHull(geom.expr, concavity.cast("double").expr, allowHoles.expr, expressionConfig))
def st_concavehull(geom: Column, concavity: Double, allowHoles: Boolean = false): Column =
ColumnAdapter(ST_ConcaveHull(geom.expr, lit(concavity).cast("double").expr, lit(allowHoles).expr, expressionConfig))
def st_difference(geom1: Column, geom2: Column): Column = ColumnAdapter(ST_Difference(geom1.expr, geom2.expr, expressionConfig))
def st_distance(geom1: Column, geom2: Column): Column = ColumnAdapter(ST_Distance(geom1.expr, geom2.expr, expressionConfig))
def st_dimension(geom: Column): Column = ColumnAdapter(ST_Dimension(geom.expr, expressionConfig))
Expand Down
Loading

0 comments on commit 8fa5d6a

Please sign in to comment.