Exposed spatial_index resolution parameter.

Author: metasim

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

@@ -59,11 +59,21 @@ trait RasterFunctions {
   /** Extracts the bounding box from a RasterSource or ProjectedRasterTile */
   def rf_extent(col: Column): TypedColumn[Any, Extent] = GetExtent(col)
 
-  /** Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS */
-  def rf_spatial_index(targetExtent: Column, targetCRS: Column) = XZ2Indexer(targetExtent, targetCRS)
+  /** Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS
+    * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html */
+  def rf_spatial_index(targetExtent: Column, targetCRS: Column, indexResolution: Short) = XZ2Indexer(targetExtent, targetCRS, indexResolution)
 
-  /** Constructs a XZ2 index in WGS84 from either a ProjectedRasterTile or RasterSource */
-  def rf_spatial_index(targetExtent: Column) = XZ2Indexer(targetExtent)
+  /** Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS
+    * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html */
+  def rf_spatial_index(targetExtent: Column, targetCRS: Column) = XZ2Indexer(targetExtent, targetCRS, 18: Short)
+
+  /** Constructs a XZ2 index with level 18 resolution in WGS84 from either a ProjectedRasterTile or RasterSource
+    * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html  */
+  def rf_spatial_index(targetExtent: Column, indexResolution: Short) = XZ2Indexer(targetExtent, indexResolution)
+
+  /** Constructs a XZ2 index with level 18 resolution in WGS84 from either a ProjectedRasterTile or RasterSource
+    * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html  */
+  def rf_spatial_index(targetExtent: Column) = XZ2Indexer(targetExtent, 18: Short)
 
   /** Extracts the CRS from a RasterSource or ProjectedRasterTile */
   def rf_crs(col: Column): TypedColumn[Any, CRS] = GetCRS(col)

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

@@ -44,7 +44,7 @@ import org.locationtech.rasterframes.expressions.accessors.GetCRS
 
 /**
   * Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource
-  *
+  * See: https://www.geomesa.org/documentation/user/datastores/index_overview.html
   * @param left geometry-like column
   * @param right CRS column
   * @param indexResolution resolution level of the space filling curve -
@@ -59,7 +59,7 @@ import org.locationtech.rasterframes.expressions.accessors.GetCRS
     * crs - the native CRS of the `geom` column
 """
 )
-case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Short = 18)
+case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Short)
   extends BinaryExpression with CodegenFallback {
 
   override def nodeName: String = "rf_spatial_index"
@@ -75,6 +75,7 @@ case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Shor
   }
 
   private lazy val indexer = XZ2SFC(indexResolution)
+  @transient
   private lazy val gf = new GeometryFactory()
 
   override protected def nullSafeEval(leftInput: Any, rightInput: Any): Any = {
@@ -121,8 +122,10 @@ case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Shor
 
 object XZ2Indexer {
   import org.locationtech.rasterframes.encoders.SparkBasicEncoders.longEnc
+  def apply(targetExtent: Column, targetCRS: Column, indexResolution: Short): TypedColumn[Any, Long] =
+    new Column(new XZ2Indexer(targetExtent.expr, targetCRS.expr, indexResolution)).as[Long]
   def apply(targetExtent: Column, targetCRS: Column): TypedColumn[Any, Long] =
-    new Column(new XZ2Indexer(targetExtent.expr, targetCRS.expr)).as[Long]
-  def apply(targetExtent: Column): TypedColumn[Any, Long] =
-    new Column(new XZ2Indexer(targetExtent.expr, GetCRS(targetExtent.expr))).as[Long]
+    new Column(new XZ2Indexer(targetExtent.expr, targetCRS.expr, 18)).as[Long]
+  def apply(targetExtent: Column, indexResolution: Short = 18): TypedColumn[Any, Long] =
+    new Column(new XZ2Indexer(targetExtent.expr, GetCRS(targetExtent.expr), indexResolution)).as[Long]
 }

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

@@ -101,11 +101,21 @@ class XZ2IndexerSpec extends TestEnvironment with Inspectors {
 
     }
     it("should work when CRS is LatLng") {
-
       val df = testExtents.map(Tuple1.apply).toDF("extent")
       val crs: CRS = LatLng
       val indexes = df.select(rf_spatial_index($"extent", serialized_literal(crs))).collect()
 
+      forEvery(indexes.zip(expected)) { case (i, e) =>
+        i should be (e)
+      }
+    }
+    it("should support custom resolution") {
+      val sfc = XZ2SFC(3)
+      val expected = testExtents.map(e => sfc.index(e.xmin, e.ymin, e.xmax, e.ymax))
+      val df = testExtents.map(Tuple1.apply).toDF("extent")
+      val crs: CRS = LatLng
+      val indexes = df.select(rf_spatial_index($"extent", serialized_literal(crs), 3)).collect()
+
       forEvery(indexes.zip(expected)) { case (i, e) =>
         i should be (e)
       }

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

@@ -586,12 +586,16 @@ def rf_geometry(proj_raster_col):
     return _apply_column_function('rf_geometry', proj_raster_col)
 
 
-def rf_spatial_index(geom_col, crs_col=None):
-    """Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS"""
+def rf_spatial_index(geom_col, crs_col=None, index_resolution = 18):
+    """Constructs a XZ2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS.
+       For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html """
+
+    jfcn = RFContext.active().lookup('rf_spatial_index')
+
     if crs_col is not None:
-        return _apply_column_function('rf_spatial_index', geom_col, crs_col)
+        return Column(jfcn(_to_java_column(geom_col), _to_java_column(crs_col), index_resolution))
     else:
-        return _apply_column_function('rf_spatial_index', geom_col)
+        return Column(jfcn(_to_java_column(geom_col), index_resolution))
 
 # ------ GeoMesa Functions ------
 

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

@@ -163,3 +163,9 @@ def test_spatial_index(self):
         indexes = {x[0] for x in df.collect()}
         self.assertSetEqual(indexes, expected)
 
+        # Custom resolution
+        df = self.df.select(rf_spatial_index(self.df.poly_geom, rf_crs(lit("EPSG:4326")), 3).alias('index'))
+        expected = {21, 36}
+        indexes = {x[0] for x in df.collect()}
+        self.assertSetEqual(indexes, expected)
+