rf_spatial_index python one arg variant; expand unit tests

Signed-off-by: Jason T. Brown <[email protected]>

Author: vpipkt

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

@@ -586,9 +586,12 @@ def rf_geometry(proj_raster_col):
     return _apply_column_function('rf_geometry', proj_raster_col)
 
 
-def rf_spatial_index(geom_col, crs_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"""
-    return _apply_column_function('rf_spatial_index', geom_col, crs_col)
+    if crs_col is not None:
+        return _apply_column_function('rf_spatial_index', geom_col, crs_col)
+    else:
+        return _apply_column_function('rf_spatial_index', geom_col)
 
 # ------ GeoMesa Functions ------
 

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

@@ -323,9 +323,6 @@ def test_render_composite(self):
         # Look for the PNG magic cookie
         self.assertEqual(png_bytes[0:8], bytearray([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]))
 
-
-
-
     def test_rf_interpret_cell_type_as(self):
         from pyspark.sql import Row
         from pyrasterframes.rf_types import Tile
@@ -359,3 +356,15 @@ def test_rf_local_data_and_no_data(self):
 
         result_d = result['ld']
         assert_equal(result_d.cells, np.invert(t.cells.mask))
+
+    def test_rf_spatial_index(self):
+        from pyspark.sql.functions import min as F_min
+        result_one_arg = self.df.select(rf_spatial_index('tile').alias('ix')) \
+                            .agg(F_min('ix')).first()[0]
+        print(result_one_arg)
+
+        result_two_arg = self.df.select(rf_spatial_index(rf_extent('tile'), rf_crs('tile')).alias('ix')) \
+                            .agg(F_min('ix')).first()[0]
+
+        self.assertEqual(result_two_arg, result_one_arg)
+        self.assertEqual(result_one_arg, 55179438768)  # this is a bit more fragile but less important

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

@@ -162,3 +162,4 @@ def test_spatial_index(self):
         expected = {22858201775, 38132946267, 38166922588, 38180072113}
         indexes = {x[0] for x in df.collect()}
         self.assertSetEqual(indexes, expected)
+

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

@@ -94,3 +94,7 @@ def create_layer(self):
         self.rf = rf.withColumn('tile2', rf_convert_cell_type('tile', 'float32')) \
             .drop('tile') \
             .withColumnRenamed('tile2', 'tile').as_layer()
+
+        df = self.spark.read.raster(self.img_uri)
+        self.df = df.withColumn('tile', rf_convert_cell_type('proj_raster', 'float32')) \
+            .drop('proj_raster')