1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, 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  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.spark.sql
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import org.apache.spark.sql.catalyst.expressions._
13 import org.apache.spark.sql.catalyst.plans.logical._
14 import org.apache.spark.sql.catalyst.rules.Rule
15 import org.apache.spark.sql.execution.datasources.LogicalRelation
16 import org.apache.spark.sql.execution.{ProjectExec, SparkPlan}
17 import org.apache.spark.sql.sedona_sql.UDT.{GeometryUDT => Sedona_GeometryUDT}
18 import org.apache.spark.sql.sedona_sql.expressions.{ST_Predicate => Sedona_ST_Predicate}
19 import org.apache.spark.sql.types.DataTypes
20 import org.apache.spark.sql.{SQLContext, Strategy}
21 import org.geotools.api.filter.expression.{Expression => GTExpression, Literal => GTLiteral}
22 import org.geotools.api.filter.{FilterFactory, Filter => GTFilter}
23 import org.geotools.factory.CommonFactoryFinder
24 import org.locationtech.geomesa.filter.FilterHelper
25 import org.locationtech.geomesa.spark.haveSedona
26 import org.locationtech.geomesa.spark.jts.rules.GeometryLiteral
27 import org.locationtech.geomesa.spark.jts.udf.SpatialRelationFunctions._
28 import org.locationtech.geomesa.spark.sql.GeoMesaRelation.PartitionedIndexedRDD
29 import org.locationtech.jts.geom.{Envelope, Geometry}
30 
31 import java.time.{LocalDateTime, ZoneId, ZoneOffset}
32 import java.util.Date
33 
34 object SQLRules extends LazyLogging {
35   @transient
36   private val ff: FilterFactory = CommonFactoryFinder.getFilterFactory
37 
38   def scalaUDFtoGTFilter(udf: Expression): Option[GTFilter] = {
39     udf match {
40       case u: ScalaUDF if u.children.length == 2 => buildGTFilter(u.function, u.children.head, u.children.last)
41       case _ => None
42     }
43   }
44 
45   private def buildGTFilter(func: AnyRef, exprA: Expression, exprB: Expression): Option[GTFilter] =
46     for {
47       builder <- funcToFF(func)
48       gtExprA <- sparkExprToGTExpr(exprA)
49       gtExprB <- sparkExprToGTExpr(exprB)
50     } yield {
51       builder(gtExprA, gtExprB)
52     }
53 
54   def funcToFF(func: AnyRef): Option[(GTExpression, GTExpression) => GTFilter] = {
55     func match {
56       case ST_Contains => Some((expr1: GTExpression, expr2: GTExpression) =>
57         ff.contains(expr1, expr2))
58       case ST_Crosses => Some((expr1: GTExpression, expr2: GTExpression) =>
59         ff.crosses(expr1, expr2))
60       case ST_Disjoint => Some((expr1: GTExpression, expr2: GTExpression) =>
61         ff.disjoint(expr1, expr2))
62       case ST_Equals => Some((expr1: GTExpression, expr2: GTExpression) =>
63         ff.equal(expr1, expr2))
64       case ST_Intersects => Some((expr1: GTExpression, expr2: GTExpression) =>
65         ff.intersects(expr1, expr2))
66       case ST_Overlaps => Some((expr1: GTExpression, expr2: GTExpression) =>
67         ff.overlaps(expr1, expr2))
68       case ST_Touches => Some((expr1: GTExpression, expr2: GTExpression) =>
69         ff.touches(expr1, expr2))
70       case ST_Within => Some((expr1: GTExpression, expr2: GTExpression) =>
71         ff.within(expr1, expr2))
72       case _ => None
73     }
74   }
75 
76   def sedonaExprToGTFilter(pred: Sedona_ST_Predicate): Option[GTFilter] = {
77     sparkExprToGTExpr(pred.children.head).flatMap { expr1 =>
78       sparkExprToGTExpr(pred.children.last).flatMap { expr2 =>
79         // sedona classes are private, so we have to match on the class name instead of the class itself
80         pred.getClass.getSimpleName match {
81           case "ST_Contains"   => Some(ff.contains(expr1, expr2))
82           case "ST_Crosses"    => Some(ff.crosses(expr1, expr2))
83           case "ST_Overlaps"   => Some(ff.overlaps(expr1, expr2))
84           case "ST_Intersects" => Some(ff.intersects(expr1, expr2))
85           case "ST_Within"     => Some(ff.within(expr1, expr2))
86           case "ST_Touches"    => Some(ff.touches(expr1, expr2))
87           case "ST_Equals"     => Some(ff.equal(expr1, expr2))
88           case "ST_Disjoint"   => Some(ff.disjoint(expr1, expr2))
89           case _ => None
90         }
91       }
92     }
93   }
94 
95   def sparkFilterToGTFilter(expr: Expression): Option[GTFilter] = {
96     expr match {
97       case udf: ScalaUDF => scalaUDFtoGTFilter(udf)
98       case binaryComp@BinaryComparison(left, right) =>
99         val leftExpr = sparkExprToGTExpr(left)
100         val rightExpr = sparkExprToGTExpr(right)
101         if (leftExpr.isEmpty || rightExpr.isEmpty) {
102           None
103         } else {
104           binaryComp match {
105             case _: EqualTo            => Some(ff.equals(leftExpr.get, rightExpr.get))
106             case _: LessThan           => Some(ff.less(leftExpr.get, rightExpr.get))
107             case _: LessThanOrEqual    => Some(ff.lessOrEqual(leftExpr.get, rightExpr.get))
108             case _: GreaterThan        => Some(ff.greater(leftExpr.get, rightExpr.get))
109             case _: GreaterThanOrEqual => Some(ff.greaterOrEqual(leftExpr.get, rightExpr.get))
110             case _ => None
111           }
112         }
113       case unary: UnaryExpression =>
114         val sparkExpr = unary.child
115         val gtExpr = sparkExprToGTExpr(sparkExpr)
116         if (gtExpr.isEmpty)
117           None
118         else {
119           unary match {
120             case _: IsNotNull => Some(ff.not(ff.isNull(gtExpr.get)))
121             case _: IsNull => Some(ff.isNull(gtExpr.get))
122             case _ => None
123           }
124         }
125       case _ =>
126         if (haveSedona && expr.isInstanceOf[Sedona_ST_Predicate]) {
127           sedonaExprToGTFilter(expr.asInstanceOf[Sedona_ST_Predicate])
128         } else {
129           logger.debug(s"Got expr: $expr.  Don't know how to turn this into a GeoTools Expression.")
130           None
131         }
132     }
133   }
134 
135   def sparkExprToGTExpr(expression: Expression): Option[GTExpression] = expression match {
136     case g: GeometryLiteral =>
137       Some(ff.literal(g.geom))
138 
139     case a: AttributeReference if a.name != "__fid__" =>
140       Some(ff.property(a.name))
141 
142     case c: Cast =>
143       lazy val zone = c.timeZoneId.map(ZoneId.of).orNull
144       sparkExprToGTExpr(c.child).map {
145         case lit: GTLiteral if lit.getValue.isInstanceOf[Date] && zone != null =>
146           val date = LocalDateTime.ofInstant(lit.getValue.asInstanceOf[Date].toInstant, zone)
147           ff.literal(new Date(date.atZone(ZoneOffset.UTC).toInstant.toEpochMilli))
148         case e => e
149       }
150 
151     case lit: Literal if lit.dataType == DataTypes.StringType =>
152       // the actual class is org.apache.spark.unsafe.types.UTF8String, we need to make it
153       // a normal string so that geotools can handle it
154       Some(ff.literal(Option(lit.value).map(_.toString).orNull))
155 
156     case lit: Literal if lit.dataType == DataTypes.TimestampType =>
157       // timestamps are defined as microseconds
158       Some(ff.literal(new Date(lit.value.asInstanceOf[Long] / 1000)))
159 
160     case lit: Literal if haveSedona && lit.dataType.isInstanceOf[Sedona_GeometryUDT] =>
161       Some(ff.literal(lit.dataType.asInstanceOf[Sedona_GeometryUDT].deserialize(lit.value)))
162 
163     case lit: Literal =>
164       Some(ff.literal(lit.value))
165 
166     case _ =>
167       logger.debug(s"Can't turn expression into geotools: $expression")
168       None
169   }
170 
171   // new optimizations rules
172   object SpatialOptimizationsRule extends Rule[LogicalPlan] with PredicateHelper {
173 
174     def extractGeometry(e: org.apache.spark.sql.catalyst.expressions.Expression): Option[Geometry] = e match {
175       case GeometryLiteral(_, geom) => Some(geom)
176       case And(l, r) => extractGeometry(l).orElse(extractGeometry(r))
177       case u: ScalaUDF => u.children.collectFirst { case GeometryLiteral(_, geom) => geom }
178       case _ => None
179     }
180 
181     private def extractGridId(envelopes: List[Envelope],
182                               e: org.apache.spark.sql.catalyst.expressions.Expression): Option[List[Int]] =
183       extractGeometry(e).map(RelationUtils.gridIdMapper(_, envelopes))
184 
185     // Replace the relation in a join with a GeoMesaJoin Relation
186     private def alterJoin(join: Join): LogicalPlan = {
187       val isSpatialUDF = join.condition.exists {
188         case u: ScalaUDF if u.function.isInstanceOf[(Geometry, Geometry) => java.lang.Boolean] =>
189           u.children.head.isInstanceOf[AttributeReference] && u.children(1).isInstanceOf[AttributeReference]
190         case _ => false
191       }
192 
193       (join.left, join.right) match {
194         case (left: LogicalRelation, right: LogicalRelation) if isSpatialUDF =>
195           (left.relation, right.relation) match {
196             case (leftRel: GeoMesaRelation, rightRel: GeoMesaRelation) =>
197               leftRel.join(rightRel, join.condition.get) match {
198                 case None => join
199                 case Some(joinRelation) =>
200                   val newLogicalRelLeft = SparkVersions.copy(left)(output = left.output ++ right.output, relation = joinRelation)
201                   SparkVersions.copy(join)(left = newLogicalRelLeft)
202               }
203 
204             case _ => join
205           }
206 
207         case (leftProject @ Project(leftProjectList, left: LogicalRelation),
208             Project(rightProjectList, right: LogicalRelation)) if isSpatialUDF =>
209           (left.relation, right.relation) match {
210             case (leftRel: GeoMesaRelation, rightRel: GeoMesaRelation) =>
211               leftRel.join(rightRel, join.condition.get) match {
212                 case None => join
213                 case Some(joinRelation) =>
214                   val newLogicalRelLeft = SparkVersions.copy(left)(output = left.output ++ right.output, relation = joinRelation)
215                   val newProjectLeft = leftProject.copy(projectList = leftProjectList ++ rightProjectList, child = newLogicalRelLeft)
216                   SparkVersions.copy(join)(left = newProjectLeft)
217               }
218 
219             case _ => join
220           }
221 
222         case _ => join
223       }
224     }
225 
226     override def apply(plan: LogicalPlan): LogicalPlan = {
227       logger.debug(s"Optimizer sees $plan")
228 
229       // NOTE: The number of arguments in Aggregate constructor is 3 on community spark and 4 on DataBricks
230       // so we cannot use pattern matching to unapply the constructor also
231       // need to use reflection to safely create new Aggregate instance
232       val optimizeAggregate: PartialFunction[LogicalPlan, LogicalPlan] =
233         Function.unlift { plan: LogicalPlan =>
234           plan match {
235             case agg: Aggregate =>
236               agg.child match {
237                 case Project(projectList, join: Join) =>
238                   val alteredJoin = SpatialOptimizationsRule.alterJoin(join)
239                   Some(Aggregates.instance(agg.groupingExpressions, agg.aggregateExpressions, Project(projectList, alteredJoin), None))
240                 case join: Join =>
241                   val alteredJoin = SpatialOptimizationsRule.alterJoin(join)
242                   Some(Aggregates.instance(agg.groupingExpressions, agg.aggregateExpressions, alteredJoin, None))
243                 case _ => None
244               }
245             case _ => None
246           }
247         }
248 
249       val optimizeRest: PartialFunction[LogicalPlan, LogicalPlan] = {
250         case join: Join =>
251           alterJoin(join)
252         case sort @ Sort(_, _, _) => sort    // No-op.  Just realizing what we can do:)
253         case filt @ Filter(f, lr: LogicalRelation) if lr.relation.isInstanceOf[GeoMesaRelation] =>
254           // TODO: deal with `or`
255 
256           val gmRel = lr.relation.asInstanceOf[GeoMesaRelation]
257           // split up conjunctive predicates and extract the st_contains variable
258           val sparkFilters: Seq[Expression] =  splitConjunctivePredicates(f)
259 
260           val (gtFilters: Seq[GTFilter], sFilters: Seq[Expression]) = sparkFilters.foldLeft((Seq[GTFilter](), Seq[Expression]())) {
261             case ((gts: Seq[GTFilter], sfilters), expression: Expression) =>
262               sparkFilterToGTFilter(expression) match {
263                 case Some(gtf) => (gts.+:(gtf), sfilters)
264                 case None      => (gts,         sfilters.+:(expression))
265               }
266           }
267 
268           if (gtFilters.nonEmpty) {
269             // if we have a partitioned cache, exclude partitions that don't match the query filter
270             val partitioned = gmRel.cached.map {
271               case c: PartitionedIndexedRDD =>
272                 val hints = sparkFilters.flatMap(extractGridId(c.envelopes, _)).flatten
273                 if (hints.isEmpty) { c } else {
274                   c.copy(rdd = c.rdd.filter { case (key, _) => hints.contains(key) })
275                 }
276 
277               case c => c
278             }
279             val filt = FilterHelper.filterListAsAnd(gmRel.filter.toSeq ++ gtFilters)
280             val relation = gmRel.copy(filter = filt, cached = partitioned)
281             val newrel = SparkVersions.copy(lr)(output = lr.output, relation = relation)
282             if (sFilters.nonEmpty) {
283               // Keep filters that couldn't be transformed at the top level
284               Filter(sFilters.reduce(And), newrel)
285             } else {
286               // if all filters could be transformed to GeoTools filters, just return the new relation
287               newrel
288             }
289           } else {
290             filt
291           }
292       }
293       plan.transform(optimizeAggregate orElse optimizeRest)
294     }
295 
296   }
297 
298   // A catch for when we are able to precompute the join using the sweepline algorithm.
299   // Skips doing a full cartesian product with catalyst.
300   object SpatialJoinStrategy extends Strategy {
301 
302     import org.apache.spark.sql.catalyst.plans.logical._
303 
304     def alterJoin(logicalPlan: Join): Seq[SparkPlan] = {
305       logicalPlan.left match {
306         case Project(projectList, lr: LogicalRelation) if lr.relation.isInstanceOf[GeoMesaJoinRelation] =>
307            ProjectExec(projectList, planLater(lr)) :: Nil
308 
309         case lr: LogicalRelation if lr.relation.isInstanceOf[GeoMesaJoinRelation] => planLater(lr) :: Nil
310 
311         case _ => Nil
312       }
313     }
314 
315     override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
316       //TODO: handle other kinds of joins
317       case Project(_, logicalPlan: Join) => alterJoin(logicalPlan)
318       case join: Join => alterJoin(join)
319       case _ => Nil
320     }
321   }
322 
323   def registerOptimizations(sqlContext: SQLContext): Unit = {
324 
325     Seq(SpatialOptimizationsRule).foreach { r =>
326       if(!sqlContext.experimental.extraOptimizations.contains(r))
327         sqlContext.experimental.extraOptimizations ++= Seq(r)
328     }
329 
330     Seq(SpatialJoinStrategy).foreach { s =>
331       if(!sqlContext.experimental.extraStrategies.contains(s))
332         sqlContext.experimental.extraStrategies ++= Seq(s)
333     }
334   }
335 }