Fix for #72.

Misc refactoring of Python RFContext methods.

Signed-off-by: Simeon H.K. Fitch <[email protected]>

Author: metasim

pyrasterframes/build.sbt

@@ -153,6 +153,7 @@ ivyPaths in pysparkCmd := ivyPaths.value.withIvyHome(target.value / "ivy")
 pyTest := {
   val _ = assembly.value
   val s = streams.value
+  s.log.info("Running python tests...")
   val wd = pythonSource.value
   Process("python setup.py test", wd) ! s.log
 }

pyrasterframes/python/geomesa_pyspark/spark.py

@@ -12,12 +12,11 @@
 from pyspark.sql.types import UserDefinedType
 from pyspark.sql import Row
 from pyspark.sql.types import *
-from pyrasterframes.context import _checked_context
+from pyrasterframes.context import RFContext
 
 
 __all__ = ['GeometryUDT']
 
-
 class GeometryUDT(UserDefinedType):
     """User-defined type (UDT).
 
@@ -41,4 +40,4 @@ def serialize(self, obj):
         return Row(obj.toBytes)
 
     def deserialize(self, datum):
-        return _checked_context().generateGeometry(datum[0])
+        return RFContext._jvm_mirror().generateGeometry(datum[0])

pyrasterframes/python/pyrasterframes/__init__.py

@@ -3,18 +3,15 @@
 appended to PySpark classes.
 """
 
-
 from __future__ import absolute_import
-from pyspark.sql.types import UserDefinedType
 from pyspark import SparkContext
 from pyspark.sql import SparkSession, DataFrame, DataFrameReader
-from pyspark.sql.types import *
 from pyspark.sql.column import _to_java_column
 
 # Import RasterFrame types and functions
 from .types import *
 from . import rasterfunctions
-
+from .context import RFContext
 
 __all__ = ['RasterFrame', 'TileExploder']
 

pyrasterframes/python/pyrasterframes/context.py

@@ -4,10 +4,42 @@
 
 from pyspark import SparkContext
 
-def _checked_context():
-    """ Get the active SparkContext and throw an error if it is not enabled for RasterFrames."""
-    sc = SparkContext._active_spark_context
-    if not hasattr(sc, '_rf_context'):
-        raise AttributeError(
-            "RasterFrames have not been enabled for the active session. Call 'SparkSession.withRasterFrames()'.")
-    return sc._rf_context._jrfctx
+__all__ = ['RFContext']
+
+
+class RFContext(object):
+    """
+    Entrypoint to RasterFrames services
+    """
+    def __init__(self, spark_session):
+        self._spark_session = spark_session
+        self._gateway = spark_session.sparkContext._gateway
+        self._jvm = self._gateway.jvm
+        jsess = self._spark_session._jsparkSession
+        self._jrfctx = self._jvm.astraea.spark.rasterframes.py.PyRFContext(jsess)
+
+    def list_to_seq(self, py_list):
+        conv = self.lookup('listToSeq')
+        return conv(py_list)
+
+    def lookup(self, function_name):
+        return getattr(self._jrfctx, function_name)
+
+    @staticmethod
+    def active():
+        """
+        Get the active Pythono RFContext and throw an error if it is not enabled for RasterFrames.
+        """
+        sc = SparkContext._active_spark_context
+        if not hasattr(sc, '_rf_context'):
+            raise AttributeError(
+                "RasterFrames have not been enabled for the active session. Call 'SparkSession.withRasterFrames()'.")
+        return sc._rf_context
+
+    @staticmethod
+    def _jvm_mirror():
+        """
+        Get the active Scala PyRFContext and throw an error if it is not enabled for RasterFrames.
+        """
+        return RFContext.active()._jrfctx
+

pyrasterframes/python/pyrasterframes/rasterfunctions.py

@@ -8,14 +8,14 @@
 from __future__ import absolute_import
 from pyspark.sql.types import *
 from pyspark.sql.column import Column, _to_java_column
-from .context import _checked_context
+from .context import RFContext
 
 
 THIS_MODULE = 'pyrasterframes'
 
 
 def _context_call(name, *args):
-    f = getattr(_checked_context(), name)
+    f = RFContext.active().lookup(name)
     return f(*args)
 
 
@@ -27,8 +27,7 @@ def _celltype(cellTypeStr):
 def _create_assembleTile():
     """ Create a function mapping to the Scala implementation."""
     def _(colIndex, rowIndex, cellData, numCols, numRows, cellType):
-        ctx = _checked_context()
-        jfcn = getattr(ctx, 'assembleTile')
+        jfcn = RFContext.active().lookup('assembleTile')
         return Column(jfcn(_to_java_column(colIndex), _to_java_column(rowIndex), _to_java_column(cellData), numCols, numRows, _celltype(cellType)))
     _.__name__ = 'assembleTile'
     _.__doc__ = "Create a Tile from  a column of cell data with location indices"
@@ -39,7 +38,7 @@ def _(colIndex, rowIndex, cellData, numCols, numRows, cellType):
 def _create_arrayToTile():
     """ Create a function mapping to the Scala implementation."""
     def _(arrayCol, numCols, numRows):
-        jfcn = getattr(_checked_context(), 'arrayToTile')
+        jfcn = RFContext.active().lookup('arrayToTile')
         return Column(jfcn(_to_java_column(arrayCol), numCols, numRows))
     _.__name__ = 'arrayToTile'
     _.__doc__ = "Convert array in `arrayCol` into a Tile of dimensions `numCols` and `numRows'"
@@ -50,7 +49,7 @@ def _(arrayCol, numCols, numRows):
 def _create_convertCellType():
     """ Create a function mapping to the Scala implementation."""
     def _(tileCol, cellType):
-        jfcn = getattr(_checked_context(), 'convertCellType')
+        jfcn = RFContext.active().lookup('convertCellType')
         return Column(jfcn(_to_java_column(tileCol), _celltype(cellType)))
     _.__name__ = 'convertCellType'
     _.__doc__ = "Convert the numeric type of the Tiles in `tileCol`"
@@ -61,7 +60,7 @@ def _(tileCol, cellType):
 def _create_makeConstantTile():
     """ Create a function mapping to the Scala implementation."""
     def _(value, cols, rows, cellType):
-        jfcn = getattr(_checked_context(), 'makeConstantTile')
+        jfcn = RFContext.active().lookup('makeConstantTile')
         return Column(jfcn(value, cols, rows, cellType))
     _.__name__ = 'makeConstantTile'
     _.__doc__ = "Constructor for constant tile column"
@@ -72,7 +71,7 @@ def _(value, cols, rows, cellType):
 def _create_tileZeros():
     """ Create a function mapping to the Scala implementation."""
     def _(cols, rows, cellType = 'float64'):
-        jfcn = getattr(_checked_context(), 'tileZeros')
+        jfcn = RFContext.active().lookup('tileZeros')
         return Column(jfcn(cols, rows, cellType))
     _.__name__ = 'tileZeros'
     _.__doc__ = "Create column of constant tiles of zero"
@@ -83,7 +82,7 @@ def _(cols, rows, cellType = 'float64'):
 def _create_tileOnes():
     """ Create a function mapping to the Scala implementation."""
     def _(cols, rows, cellType = 'float64'):
-        jfcn = getattr(_checked_context(), 'tileOnes')
+        jfcn = RFContext.active().lookup('tileOnes')
         return Column(jfcn(cols, rows, cellType))
     _.__name__ = 'tileOnes'
     _.__doc__ = "Create column of constant tiles of one"
@@ -94,7 +93,7 @@ def _(cols, rows, cellType = 'float64'):
 def _create_rasterize():
     """ Create a function mapping to the Scala rasterize function. """
     def _(geometryCol, boundsCol, valueCol, numCols, numRows):
-        jfcn = getattr(_checked_context(), 'rasterize')
+        jfcn = RFContext.active().lookup('rasterize')
         return Column(jfcn(_to_java_column(geometryCol), _to_java_column(boundsCol), _to_java_column(valueCol), numCols, numRows))
     _.__name__ = 'rasterize'
     _.__doc__ = 'Create a tile where cells in the grid defined by cols, rows, and bounds are filled with the given value.'
@@ -104,7 +103,7 @@ def _(geometryCol, boundsCol, valueCol, numCols, numRows):
 def _create_reproject_geometry():
     """ Create a function mapping to the Scala reprojectGeometry function. """
     def _(geometryCol, srcCRSName, dstCRSName):
-        jfcn = getattr(_checked_context(), 'reprojectGeometry')
+        jfcn = RFContext.active().lookup('reprojectGeometry')
         return Column(jfcn(_to_java_column(geometryCol), srcCRSName, dstCRSName))
     _.__name__ = 'reprojectGeometry'
     _.__doc__ = """Reproject a column of geometry given the CRS names of the source and destination.
@@ -114,6 +113,17 @@ def _(geometryCol, srcCRSName, dstCRSName):
     _.__module__ = THIS_MODULE
     return _
 
+def _create_explode_tiles():
+    """ Create a function mapping to Scala explodeTiles function """
+    def _(*args):
+        jfcn = RFContext.active().lookup('explodeTiles')
+        jcols = [_to_java_column(arg) for arg in args]
+        return Column(jfcn(RFContext.active().list_to_seq(jcols)))
+    _.__name__ = 'explodeTiles'
+    _.__doc__ = 'Create a row for each cell in Tile.'
+    _.__module__ = THIS_MODULE
+    return _
+
 _rf_unique_functions = {
     'assembleTile': _create_assembleTile(),
     'arrayToTile': _create_arrayToTile(),
@@ -123,7 +133,8 @@ def _(geometryCol, srcCRSName, dstCRSName):
     'tileOnes': _create_tileOnes(),
     'cellTypes': lambda: _context_call('cellTypes'),
     'rasterize': _create_rasterize(),
-    'reprojectGeometry': _create_reproject_geometry()
+    'reprojectGeometry': _create_reproject_geometry(),
+    'explodeTiles': _create_explode_tiles()
 }
 
 
@@ -154,7 +165,6 @@ def _(geometryCol, srcCRSName, dstCRSName):
 
 _rf_column_functions = {
     # ------- RasterFrames functions -------
-    'explodeTiles': 'Create a row for each cell in Tile.',
     'tileDimensions': 'Query the number of (cols, rows) in a Tile.',
     'envelope': 'Extracts the bounding box (envelope) of the geometry.',
     'tileToIntArray': 'Flattens Tile into an array of integers.',
@@ -280,7 +290,7 @@ def _(geometryCol, srcCRSName, dstCRSName):
 def _create_column_function(name, doc=""):
     """ Create a mapping to Scala UDF for a column function by name"""
     def _(*args):
-        jfcn = getattr(_checked_context(), name)
+        jfcn = RFContext.active().lookup(name)
         jcols = [_to_java_column(arg) for arg in args]
         return Column(jfcn(*jcols))
     _.__name__ = name
@@ -292,7 +302,7 @@ def _(*args):
 def _create_columnScalarFunction(name, doc=""):
     """ Create a mapping to Scala UDF for a (column, scalar) -> column function by name"""
     def _(col, scalar):
-        jfcn = getattr(_checked_context(), name)
+        jfcn = RFContext.active().lookup(name)
         return Column(jfcn(_to_java_column(col), scalar))
     _.__name__ = name
     _.__doc__ = doc

pyrasterframes/python/pyrasterframes/types.py

@@ -11,20 +11,9 @@ class here provides the PyRasterFrames entry point.
 from pyspark.sql.types import *
 from pyspark.ml.wrapper import JavaTransformer
 from pyspark.ml.util import JavaMLReadable, JavaMLWritable
-from .context import _checked_context
+from .context import RFContext
 
-__all__ = ['RFContext', 'RasterFrame', 'TileUDT', 'TileExploder', 'NoDataFilter']
-
-class RFContext(object):
-    """
-    Entrypoint to RasterFrames services
-    """
-    def __init__(self, spark_session):
-        self._spark_session = spark_session
-        self._gateway = spark_session.sparkContext._gateway
-        self._jvm = self._gateway.jvm
-        jsess = self._spark_session._jsparkSession
-        self._jrfctx = self._jvm.astraea.spark.rasterframes.py.PyRFContext(jsess)
+__all__ = ['RasterFrame', 'TileUDT', 'TileExploder', 'NoDataFilter']
 
 
 class RasterFrame(DataFrame):
@@ -156,7 +145,7 @@ def serialize(self, obj):
                   obj.toBytes)
 
     def deserialize(self, datum):
-        return _checked_context().generateTile(datum[0], datum[1], datum[2], datum[3])
+        return RFContext._jvm_mirror().generateTile(datum[0], datum[1], datum[2], datum[3])
 
 
 

pyrasterframes/python/tests/PyRasterFramesTests.py

@@ -158,12 +158,19 @@ def test_sql(self):
         self.assertTrue(_rounded_compare(statsRow.base, statsRow.double / 2))
         self.assertTrue(_rounded_compare(statsRow.base, statsRow.half * 2))
 
+    def test_explode(self):
+        self.rf.select('spatial_key', explodeTiles(self.tileCol)).show()
+
 
 def suite():
     functionTests = unittest.TestSuite()
     functionTests.addTest(RasterFunctionsTest('test_identify_columns'))
+    functionTests.addTest(RasterFunctionsTest('test_tile_operations'))
     functionTests.addTest(RasterFunctionsTest('test_general'))
+    functionTests.addTest(RasterFunctionsTest('test_rasterize'))
+    functionTests.addTest(RasterFunctionsTest('test_reproject'))
     functionTests.addTest(RasterFunctionsTest('test_aggregations'))
+    functionTests.addTest(RasterFunctionsTest('test_explode'))
     functionTests.addTest(RasterFunctionsTest('test_sql'))
     return functionTests
 

pyrasterframes/src/main/scala/astraea/spark/rasterframes/py/PyRFContext.scala

@@ -20,21 +20,15 @@ package astraea.spark.rasterframes.py
 
 import astraea.spark.rasterframes._
 import astraea.spark.rasterframes.util.CRSParser
-
 import com.vividsolutions.jts.geom.Geometry
-
-import geotrellis.raster.{ArrayTile, CellType, Tile, MultibandTile}
-import geotrellis.spark.{SpatialKey, SpaceTimeKey, TileLayerMetadata, MultibandTileLayerRDD, ContextRDD}
+import geotrellis.raster.{ArrayTile, CellType, MultibandTile}
 import geotrellis.spark.io._
-
-import org.locationtech.geomesa.spark.jts.util.WKBUtils
-
-import org.apache.spark.rdd.RDD
+import geotrellis.spark.{ContextRDD, MultibandTileLayerRDD, SpaceTimeKey, SpatialKey, TileLayerMetadata}
 import org.apache.spark.sql._
-
+import org.locationtech.geomesa.spark.jts.util.WKBUtils
 import spray.json._
-import astraea.spark.rasterframes.ml.NoDataFilter
 
+import scala.collection.JavaConverters._
 
 /**
  * py4j access wrapper to RasterFrame entry points.
@@ -193,4 +187,6 @@ class PyRFContext(implicit sparkSession: SparkSession) extends RasterFunctions
     val dst = CRSParser(dstName)
     reprojectGeometry(geometryCol, src, dst)
   }
+
+  def listToSeq(cols: java.util.ArrayList[AnyRef]): Seq[AnyRef] = cols.asScala
 }