Added Python bindings an finished up tests.

Author: metasim

core/src/main/scala/org/locationtech/rasterframes/expressions/package.scala

@@ -134,6 +134,8 @@ package object expressions {
     registry.registerExpression[RenderPNG.RenderCompositePNG]("rf_render_png")
     registry.registerExpression[RGBComposite]("rf_rgb_composite")
 
+    registry.registerExpression[XZ2Indexer]("rf_spatial_index")
+
     registry.registerExpression[transformers.ReprojectGeometry]("st_reproject")
   }
 }

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

@@ -26,7 +26,7 @@ import geotrellis.vector.Extent
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{TypeCheckFailure, TypeCheckSuccess}
 import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
-import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Expression}
+import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, Expression, ExpressionDescription}
 import org.apache.spark.sql.jts.JTSTypes
 import org.apache.spark.sql.rf.RasterSourceUDT
 import org.apache.spark.sql.types.{DataType, LongType}
@@ -51,6 +51,14 @@ import org.locationtech.rasterframes.expressions.accessors.GetCRS
   *                        i.e. how many times the space will be recursively quartered
   *                        1-18 is typical.
   */
+@ExpressionDescription(
+  usage = "_FUNC_(geom, crs) - Constructs a XZ2 index in WGS84/EPSG:4326",
+  arguments = """
+  Arguments:
+    * geom - Geometry or item with Geometry:  Extent, ProjectedRasterTile, or RasterSource
+    * crs - the native CRS of the `geom` column
+"""
+)
 case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Short = 18)
   extends BinaryExpression with CodegenFallback {
 

core/src/test/scala/org/locationtech/rasterframes/expressions/XZ2IndexerSpec.scala

@@ -21,15 +21,12 @@
 
 package org.locationtech.rasterframes.expressions
 import geotrellis.proj4.{CRS, LatLng, WebMercator}
-import org.locationtech.rasterframes._
-import geotrellis.vector.Extent
-import org.locationtech.rasterframes.TestEnvironment
-import org.apache.spark.sql.functions.lit
-import org.locationtech.rasterframes._
-import encoders.serialized_literal
 import geotrellis.raster.CellType
+import geotrellis.vector.Extent
 import org.apache.spark.sql.Encoders
 import org.locationtech.geomesa.curve.XZ2SFC
+import org.locationtech.rasterframes.{TestEnvironment, _}
+import org.locationtech.rasterframes.encoders.serialized_literal
 import org.locationtech.rasterframes.ref.{InMemoryRasterSource, RasterSource}
 import org.locationtech.rasterframes.tiles.ProjectedRasterTile
 import org.scalatest.Inspectors

docs/src/main/paradox/release-notes.md

@@ -5,6 +5,7 @@
 ### 0.8.3
 
 * Updated `rf_crs` to accept string columns containing CRS specifications. ([#366](https://github.com/locationtech/rasterframes/issues/366))
+* Added `rf_spatial_index` function. ([#368](https://github.com/locationtech/rasterframes/issues/368))
 
 ### 0.8.2
 

pyrasterframes/src/main/python/docs/reference.pymd

@@ -65,6 +65,15 @@ See also GeoMesa [st_envelope](https://www.geomesa.org/documentation/user/spark/
 
 Convert an extent to a Geometry. The extent likely comes from @ref:[`st_extent`](reference.md#st-extent) or @ref:[`rf_extent`](reference.md#rf-extent).
 
+
+### rf_spatial_index
+
+    Long rf_spatial_index(Geometry geom, CRS crs)
+    Long rf_spatial_index(Extent extent, CRS crs)
+    Long rf_spatial_index(ProjectedRasterTile proj_raster, CRS crs)
+    
+Constructs a XZ2 index in WGS84/EPSG:4326 from either a Geometry, Extent, ProjectedRasterTile and its CRS.
+
 ## Tile Metadata and Mutation
 
 Functions to access and change the particulars of a `tile`: its shape and the data type of its cells. See section on @ref:["NoData" handling](nodata-handling.md) for additional discussion of cell types.

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -585,6 +585,11 @@ def rf_geometry(proj_raster_col):
     """Get the extent of a RasterSource or ProjectdRasterTile as a Geometry"""
     return _apply_column_function('rf_geometry', proj_raster_col)
 
+
+def rf_spatial_index(geom_col, crs_col):
+    """Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS"""
+    return _apply_column_function('rf_spatial_index', geom_col, crs_col)
+
 # ------ GeoMesa Functions ------
 
 def st_geomFromGeoHash(*args):

pyrasterframes/src/main/python/tests/VectorTypesTests.py

@@ -156,3 +156,9 @@ def test_geojson(self):
         geo = self.spark.read.geojson(sample)
         geo.show()
         self.assertEqual(geo.select('geometry').count(), 8)
+
+    def test_spatial_index(self):
+        df = self.df.select(rf_spatial_index(self.df.poly_geom, rf_crs(lit("EPSG:4326"))).alias('index'))
+        expected = {22858201775, 38132946267, 38166922588, 38180072113}
+        indexes = {x[0] for x in df.collect()}
+        self.assertSetEqual(indexes, expected)