PR feedback.

Author: metasim

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

@@ -62,38 +62,38 @@ 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
+  /** 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_xz2_index(targetExtent: Column, targetCRS: Column, indexResolution: Short) = XZ2Indexer(targetExtent, targetCRS, indexResolution)
 
   /** 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_xz2_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
+  /** Constructs a XZ2 index with provided resolution level in WGS84 from either a ProjectedRasterTile or RasterSource.
     * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html  */
   def rf_xz2_index(targetExtent: Column, indexResolution: Short) = XZ2Indexer(targetExtent, indexResolution)
 
-  /** Constructs a XZ2 index with level 18 resolution in WGS84 from either a ProjectedRasterTile or RasterSource
+  /** 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_xz2_index(targetExtent: Column) = XZ2Indexer(targetExtent, 18: Short)
 
-  /** Constructs a Z2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS
+  /** Constructs a Z2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS.
     * First the native extent is extracted or computed, and then center is used as the indexing location.
     * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html */
   def rf_z2_index(targetExtent: Column, targetCRS: Column, indexResolution: Short) = Z2Indexer(targetExtent, targetCRS, indexResolution)
 
-  /** Constructs a Z2 index in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS
+  /** Constructs a Z2 index with index resolution of 31 in WGS84 from either a Geometry, Extent, ProjectedRasterTile, or RasterSource and its CRS.
     * First the native extent is extracted or computed, and then center is used as the indexing location.
     * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html */
   def rf_z2_index(targetExtent: Column, targetCRS: Column) = Z2Indexer(targetExtent, targetCRS, 31: Short)
 
-  /** Constructs a Z2 index with level 18 resolution in WGS84 from either a ProjectedRasterTile or RasterSource
+  /** Constructs a Z2 index with the given index resolution in WGS84 from either a ProjectedRasterTile or RasterSource
     * First the native extent is extracted or computed, and then center is used as the indexing location.
     * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html  */
   def rf_z2_index(targetExtent: Column, indexResolution: Short) = Z2Indexer(targetExtent, indexResolution)
 
-  /** Constructs a Z2 index with level 18 resolution in WGS84 from either a ProjectedRasterTile or RasterSource
+  /** Constructs a Z2 index with index resolution of 31 in WGS84 from either a ProjectedRasterTile or RasterSource
     * First the native extent is extracted or computed, and then center is used as the indexing location.
     * For details: https://www.geomesa.org/documentation/user/datastores/index_overview.html  */
   def rf_z2_index(targetExtent: Column) = Z2Indexer(targetExtent, 31: Short)

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

@@ -61,9 +61,9 @@ case class XZ2Indexer(left: Expression, right: Expression, indexResolution: Shor
 
   override def checkInputDataTypes(): TypeCheckResult = {
     if (!envelopeExtractor.isDefinedAt(left.dataType))
-      TypeCheckFailure(s"Input type '${left.dataType}' does not look like something with an Extent or something with one.")
+      TypeCheckFailure(s"Input type '${left.dataType}' does not look like a geometry, extent, or something with one.")
     else if(!crsExtractor.isDefinedAt(right.dataType))
-      TypeCheckFailure(s"Input type '${right.dataType}' does not look like something with a CRS.")
+      TypeCheckFailure(s"Input type '${right.dataType}' does not look like a CRS or something with one.")
     else TypeCheckSuccess
   }
 

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

@@ -62,9 +62,9 @@ case class Z2Indexer(left: Expression, right: Expression, indexResolution: Short
 
   override def checkInputDataTypes(): TypeCheckResult = {
     if (!centroidExtractor.isDefinedAt(left.dataType))
-      TypeCheckFailure(s"Input type '${left.dataType}' does not look like something with a point.")
+      TypeCheckFailure(s"Input type '${left.dataType}' does not look like something with a centroid.")
     else if(!crsExtractor.isDefinedAt(right.dataType))
-      TypeCheckFailure(s"Input type '${right.dataType}' does not look like something with a CRS.")
+      TypeCheckFailure(s"Input type '${right.dataType}' does not look like a CRS or something with one.")
     else TypeCheckSuccess
   }
 

datasource/src/main/scala/org/locationtech/rasterframes/datasource/raster/RasterSourceRelation.scala

@@ -42,7 +42,10 @@ import org.locationtech.rasterframes.tiles.ProjectedRasterTile
   * @param bandIndexes band indexes to fetch
   * @param subtileDims how big to tile/subdivide rasters info
   * @param lazyTiles if true, creates a lazy representation of tile instead of fetching contents.
-  * @param spatialIndexPartitions if not None, adds a spatial index. If Option value < 1, uses the value of `numShufflePartitions` in SparkContext.
+  * @param spatialIndexPartitions Number of spatial index-based partitions to create.
+  *                               If Option value > 0, that number of partitions are created after adding a spatial index.
+  *                               If Option value <= 0, uses the value of `numShufflePartitions` in SparkContext.
+  *                               If None, no spatial index is added and hash partitioning is used.
   */
 case class RasterSourceRelation(
   sqlContext: SQLContext,

docs/src/main/paradox/reference.md

@@ -81,7 +81,7 @@ Constructs a XZ2 index in WGS84/EPSG:4326 from either a Geometry, Extent, Projec
     Long rf_z2_index(Extent extent, CRS crs)
     Long rf_z2_index(ProjectedRasterTile proj_raster)
 
-Constructs a Z2 index in WGS84/EPSG:4326 from either a Geometry, Extent, ProjectedRasterTile and its CRS. First the native extent is extracted or computed, and then center is used as the indexing location. This function is useful for [range partitioning](http://spark.apache.org/docs/latest/api/python/pyspark.sql.html?highlight=registerjava#pyspark.sql.DataFrame.repartitionByRange).
+Constructs a Z2 index in WGS84/EPSG:4326 from either a Geometry, Extent, ProjectedRasterTile and its CRS. First the native extent is extracted or computed, and then center is used as the indexing location. This function is useful for [range partitioning](http://spark.apache.org/docs/latest/api/python/pyspark.sql.html?highlight=registerjava#pyspark.sql.DataFrame.repartitionByRange). See @ref:[Reading Raster Data](raster-read.md#spatial-indexing-and-partitioning) section for details on how to have an index automatically added when reading raster data.
 
 ## Tile Metadata and Mutation
 

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

@@ -23,7 +23,7 @@
 * Updated to GeoTrellis 2.3.3 and Proj4j 1.1.0.
 * Fixed issues with `LazyLogger` and shading assemblies ([#293](https://github.com/locationtech/rasterframes/issues/293))
 * Updated `rf_crs` to accept string columns containing CRS specifications. ([#366](https://github.com/locationtech/rasterframes/issues/366))
-* Added `rf_xz2_index` function. ([#368](https://github.com/locationtech/rasterframes/issues/368))
+* Added `rf_spatial_index` function. ([#368](https://github.com/locationtech/rasterframes/issues/368))
 * _Breaking_ (potentially): removed `pyrasterframes.create_spark_session` in lieu of `pyrasterframes.utils.create_rf_spark_session`
 
 ### 0.8.2

pyrasterframes/src/main/python/pyrasterframes/__init__.py

@@ -149,11 +149,15 @@ def temp_name():
     if band_indexes is None:
         band_indexes = [0]
 
-    if spatial_index_partitions is False:
-        spatial_index_partitions = None
-
-    if spatial_index_partitions is not None:
-        if spatial_index_partitions is True:
+    if spatial_index_partitions:
+        num = int(spatial_index_partitions)
+        if num < 0:
+            spatial_index_partitions = '-1'
+        elif num == 0:
+            spatial_index_partitions = None
+
+    if spatial_index_partitions:
+        if spatial_index_partitions == True:
             spatial_index_partitions = "-1"
         else:
             spatial_index_partitions = str(spatial_index_partitions)

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

@@ -213,7 +213,12 @@ def test_catalog_named_arg(self):
         self.assertTrue(df.select('b1_path').distinct().count() == 3)
 
     def test_spatial_partitioning(self):
-        df = self.spark.read.raster(self.path(1, 1), spatial_index_partitions=True)
+        f = self.path(1, 1)
+        df = self.spark.read.raster(f, spatial_index_partitions=True)
         self.assertTrue('spatial_index' in df.columns)
-        # Other tests?
 
+        self.assertEqual(df.rdd.getNumPartitions(), int(self.spark.conf.get("spark.sql.shuffle.partitions")))
+        self.assertEqual(self.spark.read.raster(f, spatial_index_partitions=34).rdd.getNumPartitions(), 34)
+        self.assertEqual(self.spark.read.raster(f, spatial_index_partitions="42").rdd.getNumPartitions(), 42)
+        self.assertFalse('spatial_index' in self.spark.read.raster(f, spatial_index_partitions=False).columns)
+        self.assertFalse('spatial_index' in self.spark.read.raster(f, spatial_index_partitions=0).columns)