Skip to content

Commit 6a2132a

Browse files
metasimmetasim
committed
Rough sketch of a spatial index generator.
1 parent d59bb56 commit 6a2132a

File tree

4 files changed

+131
-0
lines changed

4 files changed

+131
-0
lines changed

‎core/src/main/scala/org/locationtech/rasterframes/RasterFunctions.scala‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,9 @@ trait RasterFunctions {
5959
/** Extracts the bounding box from a RasterSource or ProjectedRasterTile */
6060
def rf_extent(col: Column): TypedColumn[Any, Extent] = GetExtent(col)
6161

62+
/** Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource */
63+
def rf_spatial_index(targetExtent: Column, targetCRS: Column) = XZ2Indexer(targetExtent, targetCRS)
64+
6265
/** Extracts the CRS from a RasterSource or ProjectedRasterTile */
6366
def rf_crs(col: Column): TypedColumn[Any, CRS] = GetCRS(col)
6467

‎core/src/main/scala/org/locationtech/rasterframes/expressions/DynamicExtractors.scala‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,14 @@ package org.locationtech.rasterframes.expressions
2323

2424
import geotrellis.proj4.CRS
2525
import geotrellis.raster.{CellGrid, Tile}
26+
import geotrellis.vector.Extent
2627
import org.apache.spark.sql.Row
2728
import org.apache.spark.sql.catalyst.InternalRow
29+
import org.apache.spark.sql.jts.JTSTypes
2830
import org.apache.spark.sql.rf.{RasterSourceUDT, TileUDT}
2931
import org.apache.spark.sql.types._
3032
import org.apache.spark.unsafe.types.UTF8String
33+
import org.locationtech.jts.geom.Envelope
3134
import org.locationtech.rasterframes.encoders.CatalystSerializer._
3235
import org.locationtech.rasterframes.model.{LazyCRS, TileContext}
3336
import org.locationtech.rasterframes.ref.{ProjectedRasterLike, RasterRef, RasterSource}
@@ -94,6 +97,15 @@ object DynamicExtractors {
9497
(v: Any) => v.asInstanceOf[InternalRow].to[CRS]
9598
}
9699

100+
lazy val extentLikeExtractor: PartialFunction[DataType, Any ⇒ Extent] = {
101+
case t if org.apache.spark.sql.rf.WithTypeConformity(t).conformsTo(JTSTypes.GeometryTypeInstance) =>
102+
(input: Any) => JTSTypes.GeometryTypeInstance.deserialize(input).getEnvelopeInternal
103+
case t if t.conformsTo[Extent] =>
104+
(input: Any) => input.asInstanceOf[InternalRow].to[Extent]
105+
case t if t.conformsTo[Envelope] =>
106+
(input: Any) => Extent(input.asInstanceOf[InternalRow].to[Envelope])
107+
}
108+
97109
sealed trait TileOrNumberArg
98110
sealed trait NumberArg extends TileOrNumberArg
99111
case class TileArg(tile: Tile, ctx: Option[TileContext]) extends TileOrNumberArg

‎core/src/main/scala/org/locationtech/rasterframes/expressions/transformers/XZ2Indexer.scala‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
/*
2+
* This software is licensed under the Apache 2 license, quoted below.
3+
*
4+
* Copyright 2019 Astraea, Inc.
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
7+
* use this file except in compliance with the License. You may obtain a copy of
8+
* the License at
9+
*
10+
* [http://www.apache.org/licenses/LICENSE-2.0]
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
14+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
15+
* License for the specific language governing permissions and limitations under
16+
* the License.
17+
*
18+
* SPDX-License-Identifier: Apache-2.0
19+
*
20+
*/
21+
22+
package org.locationtech.rasterframes.expressions.transformers
23+
24+
import geotrellis.proj4.LatLng
25+
import geotrellis.vector.Extent
26+
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
27+
import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
28+
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
29+
import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Expression}
30+
import org.apache.spark.sql.jts.JTSTypes
31+
import org.apache.spark.sql.rf.RasterSourceUDT
32+
import org.apache.spark.sql.types.{DataType, LongType}
33+
import org.apache.spark.sql.{Column, TypedColumn}
34+
import org.locationtech.geomesa.curve.XZ2SFC
35+
import org.locationtech.jts.geom.{Envelope, Geometry, GeometryFactory}
36+
import org.locationtech.rasterframes.encoders.CatalystSerializer._
37+
import org.locationtech.rasterframes.expressions.DynamicExtractors._
38+
import org.locationtech.rasterframes.expressions.row
39+
import org.locationtech.rasterframes.jts.ReprojectionTransformer
40+
import org.locationtech.rasterframes.ref.{RasterRef, RasterSource}
41+
import org.locationtech.rasterframes.tiles.ProjectedRasterTile
42+
import org.apache.spark.sql.rf
43+
44+
/**
45+
* This expression constructs a XZ2 index for a given JTS geometry.
46+
*
47+
* @param indexResolution resolution level of the space filling curve -
48+
* i.e. how many times the space will be recursively quartered
49+
* 1-18 is typical.
50+
*/
51+
case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Short = 18)
52+
extends BinaryExpression with CodegenFallback {
53+
54+
override def nodeName: String = "rf_spatial_index"
55+
56+
override def dataType: DataType = LongType
57+
58+
override def checkInputDataTypes(): TypeCheckResult = {
59+
if (!extentLikeExtractor.orElse(projectedRasterLikeExtractor).isDefinedAt(left.dataType))
60+
TypeCheckFailure(s"Input type '${left.dataType}' does not look like something with an Extent or something with one.")
61+
else if(!crsExtractor.isDefinedAt(right.dataType))
62+
TypeCheckFailure(s"Input type '${right.dataType}' does not look like something with a CRS.")
63+
else TypeCheckSuccess
64+
}
65+
66+
private lazy val indexer = XZ2SFC(indexResolution)
67+
private lazy val gf = new GeometryFactory()
68+
69+
override protected def nullSafeEval(leftInput: Any, rightInput: Any): Any = {
70+
val crs = crsExtractor(right.dataType)(rightInput)
71+
72+
val coords = left.dataType match {
73+
case t if rf.WithTypeConformity(t).conformsTo(JTSTypes.GeometryTypeInstance) =>
74+
JTSTypes.GeometryTypeInstance.deserialize(left)
75+
case t if t.conformsTo[Extent] =>
76+
row(leftInput).to[Extent]
77+
case t if t.conformsTo[Envelope] =>
78+
row(leftInput).to[Envelope]
79+
case _: RasterSourceUDT ⇒
80+
row(leftInput).to[RasterSource](RasterSourceUDT.rasterSourceSerializer).extent
81+
case t if t.conformsTo[ProjectedRasterTile] =>
82+
row(leftInput).to[ProjectedRasterTile].extent
83+
case t if t.conformsTo[RasterRef] =>
84+
row(leftInput).to[RasterRef].extent
85+
}
86+
87+
// If no transformation is needed then just normalize to an Envelope
88+
val env = if(crs == LatLng) coords match {
89+
case e: Extent => e.jtsEnvelope
90+
case g: Geometry => g.getEnvelopeInternal
91+
case e: Envelope => e
92+
}
93+
// Otherwise convert to geometry, transform, and get envelope
94+
else {
95+
val trans = new ReprojectionTransformer(crs, LatLng)
96+
coords match {
97+
case e: Extent => trans(e.jtsGeom).getEnvelopeInternal
98+
case g: Geometry => trans(g).getEnvelopeInternal
99+
case e: Envelope => trans(gf.toGeometry(e)).getEnvelopeInternal
100+
}
101+
}
102+
103+
val index = indexer.index(
104+
env.getMinX, env.getMinY, env.getMaxX, env.getMaxY,
105+
lenient = false
106+
)
107+
index
108+
}
109+
}
110+
111+
object XZ2Indexer {
112+
import org.locationtech.rasterframes.encoders.SparkBasicEncoders.longEnc
113+
def apply(targetExtent: Column, targetCRS: Column): TypedColumn[Any, Long] =
114+
new Column(new XZ2Indexer(targetExtent.expr, targetCRS.expr)).as[Long]
115+
}

‎core/src/main/scala/org/locationtech/rasterframes/jts/ReprojectionTransformer.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import geotrellis.proj4.CRS
3131
* @since 6/4/18
3232
*/
3333
class ReprojectionTransformer(src: CRS, dst: CRS) extends GeometryTransformer {
34+
def apply(geometry: Geometry): Geometry = transform(geometry)
3435
lazy val transform = geotrellis.proj4.Transform(src, dst)
3536
override def transformCoordinates(coords: CoordinateSequence, parent: Geometry): CoordinateSequence = {
3637
val fact = parent.getFactory

0 commit comments

Comments
 (0)