1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.spark.sql
10 
11 import org.apache.spark.sql.SQLContext
12 import org.geotools.geometry.jts.GeometryCoordinateSequenceTransformer
13 import org.geotools.referencing.crs.DefaultGeographicCRS
14 import org.geotools.referencing.{CRS, GeodeticCalculator}
15 import org.locationtech.geomesa.spark.jts.util.SQLFunctionHelper.nullableUDF
16 import org.locationtech.jts.geom.{Coordinate, Geometry, LineString}
17 
18 object GeometricDistanceFunctions {
19   import java.{lang => jl}
20 
21   val ST_DistanceSpheroid: (Geometry, Geometry) => jl.Double =
22     nullableUDF((s, e) => fastDistance(s.getCoordinate, e.getCoordinate))
23 
24   // Assumes input is two points, for use with collect_list and window functions
25   val ST_AggregateDistanceSpheroid: Seq[Geometry] => jl.Double = a => ST_DistanceSpheroid(a(0), a(1))
26 
27   val ST_LengthSpheroid: LineString => jl.Double =
28     nullableUDF(line => line.getCoordinates.sliding(2).map { case Array(l, r) => fastDistance(l, r) }.sum)
29 
30   val ST_Transform: (Geometry, String, String) => Geometry = nullableUDF { (geometry, fromCRSCode, toCRSCode) =>
31     val transformer = new GeometryCoordinateSequenceTransformer
32     val fromCode = CRS.decode(fromCRSCode, true)
33     val toCode = CRS.decode(toCRSCode, true)
34     transformer.setMathTransform(CRS.findMathTransform(fromCode, toCode, true))
35     transformer.transform(geometry)
36   }
37 
38   val distanceNames = Map(
39     ST_DistanceSpheroid -> "st_distanceSpheroid",
40     ST_AggregateDistanceSpheroid -> "st_aggregateDistanceSpheroid",
41     ST_LengthSpheroid -> "st_lengthSpheroid",
42     ST_Transform -> "st_transform"
43   )
44 
45 
46   def registerFunctions(sqlContext: SQLContext): Unit = {
47     sqlContext.udf.register(distanceNames (ST_DistanceSpheroid), ST_DistanceSpheroid)
48     sqlContext.udf.register(distanceNames (ST_AggregateDistanceSpheroid), ST_AggregateDistanceSpheroid)
49     sqlContext.udf.register(distanceNames (ST_LengthSpheroid), ST_LengthSpheroid)
50     sqlContext.udf.register(distanceNames (ST_Transform), ST_Transform)
51   }
52 
53   @transient private val geoCalcs = new ThreadLocal[GeodeticCalculator] {
54     override def initialValue(): GeodeticCalculator = new GeodeticCalculator(DefaultGeographicCRS.WGS84)
55   }
56 
57   def fastDistance(c1: Coordinate, c2: Coordinate): Double = {
58     val calc = geoCalcs.get()
59     calc.setStartingGeographicPoint(c1.x, c1.y)
60     calc.setDestinationGeographicPoint(c2.x, c2.y)
61     calc.getOrthodromicDistance
62   }
63 
64 
65 }