Merge pull request #636 from jbouffard/bug-fix/reproject-partitioner

Reproject Methods Can Now Use the Specified Partitioner
locationtech-labs · Mar 9, 2018 · 48f2d3c · 48f2d3c
2 parents ae80917 + 9d463a7
commit 48f2d3c
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 28 deletions.
diff --git a/...yspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/ProjectedRasterLayer.scala b/...yspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/ProjectedRasterLayer.scala
@@ -111,22 +111,24 @@ class ProjectedRasterLayer(val rdd: RDD[(ProjectedExtent, MultibandTile)]) exten
     resampleMethod: ResampleMethod,
     partitionStrategy: PartitionStrategy
   ): TiledRasterLayer[SpatialKey] = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
     val crs = TileLayer.getCRS(targetCRS).get
     val tiled = tileToLayout(LocalLayout(256), resampleMethod, partitionStrategy).rdd
 
     layoutType match {
       case GlobalLayout(tileSize, null, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (zoom, reprojected) = tiled.reproject(crs, scheme, resampleMethod)
+        val (zoom, reprojected) = tiled.reproject(crs, scheme, resampleMethod, partitioner)
         new SpatialTiledRasterLayer(Some(zoom), reprojected)
 
       case GlobalLayout(tileSize, zoom, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (_, reprojected) = tiled.reproject(crs, scheme.levelForZoom(zoom).layout, resampleMethod)
+        val (_, reprojected) = TileRDDReproject(tiled, crs, Right(scheme.levelForZoom(zoom).layout), resampleMethod, partitioner)
         new SpatialTiledRasterLayer(Some(zoom), reprojected)
 
       case LocalLayout(tileCols, tileRows) =>
-        val (_, reprojected) = tiled.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod)
+        val (_, reprojected) = tiled.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod, partitioner)
         new SpatialTiledRasterLayer(None, reprojected)
     }
   }
@@ -137,8 +139,11 @@ class ProjectedRasterLayer(val rdd: RDD[(ProjectedExtent, MultibandTile)]) exten
     resampleMethod: ResampleMethod,
     partitionStrategy: PartitionStrategy
   ): TiledRasterLayer[SpatialKey] = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
     val tiled = tileToLayout(layoutDefinition, resampleMethod, partitionStrategy).rdd
-    val (zoom, reprojected) = TileRDDReproject(tiled, TileLayer.getCRS(target_crs).get, Right(layoutDefinition), resampleMethod)
+    val (zoom, reprojected) =
+      TileRDDReproject(tiled, TileLayer.getCRS(target_crs).get, Right(layoutDefinition), resampleMethod, partitioner)
 
     SpatialTiledRasterLayer(Some(zoom), reprojected)
   }

diff --git a/...ark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/SpatialTiledRasterLayer.scala b/...ark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/SpatialTiledRasterLayer.scala
@@ -68,34 +68,54 @@ class SpatialTiledRasterLayer(
     PythonTranslator.toPython[MultibandTile, ProtoMultibandTile](tiles)
   }
 
-  def reproject(targetCRS: String, resampleMethod: ResampleMethod): SpatialTiledRasterLayer = {
+  def reproject(
+    targetCRS: String,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): SpatialTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
     val crs = TileLayer.getCRS(targetCRS).get
     val targetLayout = FloatingLayoutScheme(rdd.metadata.layout.tileCols, rdd.metadata.layout.tileRows)
-    val (zoom, reprojected) = rdd.reproject(crs, targetLayout, resampleMethod)
+    val (zoom, reprojected) = rdd.reproject(crs, targetLayout, resampleMethod, partitioner)
     new SpatialTiledRasterLayer(Some(zoom), reprojected)
   }
 
-  def reproject(targetCRS: String, layoutType: LayoutType, resampleMethod: ResampleMethod): SpatialTiledRasterLayer = {
+  def reproject(
+    targetCRS: String,
+    layoutType: LayoutType,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): SpatialTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
     val crs = TileLayer.getCRS(targetCRS).get
+
     layoutType match {
       case GlobalLayout(tileSize, null, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (zoom, reprojected) = rdd.reproject(crs, scheme, resampleMethod)
+        val (zoom, reprojected) = rdd.reproject(crs, scheme, resampleMethod, partitioner)
         SpatialTiledRasterLayer(Some(zoom), reprojected)
 
       case GlobalLayout(tileSize, zoom, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (_, reprojected) = rdd.reproject(crs, scheme.levelForZoom(zoom).layout, resampleMethod)
+        val (_, reprojected) = TileRDDReproject(rdd, crs, Right(scheme.levelForZoom(zoom).layout), resampleMethod, partitioner)
         SpatialTiledRasterLayer(Some(zoom), reprojected)
 
       case LocalLayout(tileCols, tileRows) =>
-        val (_, reprojected) = rdd.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod)
+        val (_, reprojected) = rdd.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod, partitioner)
         SpatialTiledRasterLayer(None, reprojected)
     }
   }
 
-  def reproject(targetCRS: String, layoutDefinition: LayoutDefinition, resampleMethod: ResampleMethod): SpatialTiledRasterLayer = {
-    val (zoom, reprojected) = TileRDDReproject(rdd, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod)
+  def reproject(
+    targetCRS: String,
+    layoutDefinition: LayoutDefinition,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): SpatialTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
+    val (zoom, reprojected) = TileRDDReproject(rdd, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod, partitioner)
     SpatialTiledRasterLayer(Some(zoom), reprojected)
   }
 

diff --git a/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TemporalRasterLayer.scala b/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TemporalRasterLayer.scala
@@ -111,22 +111,24 @@ class TemporalRasterLayer(val rdd: RDD[(TemporalProjectedExtent, MultibandTile)]
     resampleMethod: ResampleMethod,
     partitionStrategy: PartitionStrategy
   ): TiledRasterLayer[SpaceTimeKey] = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
     val crs = TileLayer.getCRS(targetCRS).get
     val tiled = tileToLayout(LocalLayout(256), resampleMethod, partitionStrategy).rdd
 
     layoutType match {
       case GlobalLayout(tileSize, null, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (zoom, reprojected) = tiled.reproject(crs, scheme, resampleMethod)
+        val (zoom, reprojected) = tiled.reproject(crs, scheme, resampleMethod, partitioner)
         new TemporalTiledRasterLayer(Some(zoom), reprojected)
 
       case GlobalLayout(tileSize, zoom, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (_, reprojected) = tiled.reproject(crs, scheme.levelForZoom(zoom).layout, resampleMethod)
+        val (_, reprojected) = TileRDDReproject(tiled, crs, Right(scheme.levelForZoom(zoom).layout), resampleMethod, partitioner)
         new TemporalTiledRasterLayer(Some(zoom), reprojected)
 
       case LocalLayout(tileCols, tileRows) =>
-        val (_, reprojected) = tiled.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod)
+        val (_, reprojected) = tiled.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod, partitioner)
         new TemporalTiledRasterLayer(None, reprojected)
     }
   }
@@ -137,8 +139,10 @@ class TemporalRasterLayer(val rdd: RDD[(TemporalProjectedExtent, MultibandTile)]
     resampleMethod: ResampleMethod,
     partitionStrategy: PartitionStrategy
   ): TiledRasterLayer[SpaceTimeKey] = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
     val tiled = tileToLayout(layoutDefinition, resampleMethod, partitionStrategy).rdd
-    val (zoom, reprojected) = TileRDDReproject(tiled, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod)
+    val (zoom, reprojected) =
+      TileRDDReproject(tiled, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod, partitioner)
 
     TemporalTiledRasterLayer(Some(zoom), reprojected)
   }

diff --git a/...rk-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TemporalTiledRasterLayer.scala b/...rk-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TemporalTiledRasterLayer.scala
@@ -191,34 +191,53 @@ class TemporalTiledRasterLayer(
       .sortWith({ (t1, t2) => (t1._1.compareTo(t2._1) <= 0) })
   }
 
-  def reproject(targetCRS: String, resampleMethod: ResampleMethod): TemporalTiledRasterLayer = {
+  def reproject(
+    targetCRS: String,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TemporalTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
     val crs = TileLayer.getCRS(targetCRS).get
     val targetLayout = FloatingLayoutScheme(rdd.metadata.layout.tileCols, rdd.metadata.layout.tileRows)
-    val (zoom, reprojected) = rdd.reproject(crs, targetLayout, resampleMethod)
+    val (zoom, reprojected) = rdd.reproject(crs, targetLayout, resampleMethod, partitioner)
     TemporalTiledRasterLayer(Some(zoom), reprojected)
   }
 
-  def reproject(targetCRS: String, layoutType: LayoutType, resampleMethod: ResampleMethod): TemporalTiledRasterLayer = {
+  def reproject(
+    targetCRS: String,
+    layoutType: LayoutType,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TemporalTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
     val crs = TileLayer.getCRS(targetCRS).get
+
     layoutType match {
       case GlobalLayout(tileSize, null, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (zoom, reprojected) = rdd.reproject(crs, scheme, resampleMethod)
+        val (zoom, reprojected) = rdd.reproject(crs, scheme, resampleMethod, partitioner)
         TemporalTiledRasterLayer(Some(zoom), reprojected)
 
       case GlobalLayout(tileSize, zoom, threshold) =>
         val scheme = new ZoomedLayoutScheme(crs, tileSize, threshold)
-        val (_, reprojected) = rdd.reproject(crs, scheme.levelForZoom(zoom).layout, resampleMethod)
+        val (_, reprojected) = TileRDDReproject(rdd, crs, Right(scheme.levelForZoom(zoom).layout), resampleMethod, partitioner)
         TemporalTiledRasterLayer(Some(zoom), reprojected)
 
       case LocalLayout(tileCols, tileRows) =>
-        val (_, reprojected) = rdd.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod)
+        val (_, reprojected) = rdd.reproject(crs, FloatingLayoutScheme(tileCols, tileRows), resampleMethod, partitioner)
         TemporalTiledRasterLayer(None, reprojected)
     }
   }
 
-  def reproject(targetCRS: String, layoutDefinition: LayoutDefinition, resampleMethod: ResampleMethod): TemporalTiledRasterLayer = {
-    val (zoom, reprojected) = TileRDDReproject(rdd, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod)
+  def reproject(
+    targetCRS: String,
+    layoutDefinition: LayoutDefinition,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TemporalTiledRasterLayer = {
+    val partitioner = TileLayer.getPartitioner(partitionStrategy, rdd.getNumPartitions)
+
+    val (zoom, reprojected) = TileRDDReproject(rdd, TileLayer.getCRS(targetCRS).get, Right(layoutDefinition), resampleMethod, partitioner)
     TemporalTiledRasterLayer(Some(zoom), reprojected)
   }
 

diff --git a/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TileLayer.scala b/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TileLayer.scala
@@ -249,6 +249,12 @@ object TileLayer {
     }
   }
 
+  def getPartitioner(partitionStrategy: PartitionStrategy, defaultNumPartitions: Int): Option[Partitioner] =
+    partitionStrategy match {
+      case ps: PartitionStrategy => ps.producePartitioner(defaultNumPartitions)
+      case null => None
+    }
+
   def getCRS(crs: String): Option[CRS] = {
     Option(crs).flatMap { crs =>
       Try(CRS.fromName(crs))

diff --git a/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TiledRasterLayer.scala b/geopyspark-backend/geotrellis/src/main/scala/geopyspark/geotrellis/TiledRasterLayer.scala
@@ -124,9 +124,25 @@ abstract class TiledRasterLayer[K: SpatialComponent: JsonFormat: ClassTag: Bound
 
   protected def mask(geometries: Seq[MultiPolygon]): TiledRasterLayer[K]
 
-  protected def reproject(target_crs: String, resampleMethod: ResampleMethod): TiledRasterLayer[K]
-  protected def reproject(target_crs: String, layoutType: LayoutType, resampleMethod: ResampleMethod): TiledRasterLayer[K]
-  def reproject(targetCRS: String, layoutDefinition: LayoutDefinition, resampleMethod: ResampleMethod): TiledRasterLayer[K]
+  def reproject(
+    targetCRS: String,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TiledRasterLayer[K]
+
+  def reproject(
+    targetCRS: String,
+    layoutType: LayoutType,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TiledRasterLayer[K]
+
+  def reproject(
+    targetCRS: String,
+    layoutDefinition: LayoutDefinition,
+    resampleMethod: ResampleMethod,
+    partitionStrategy: PartitionStrategy
+  ): TiledRasterLayer[K]
 
   def tileToLayout(
     layoutDefinition: LayoutDefinition,

diff --git a/geopyspark/geotrellis/layer.py b/geopyspark/geotrellis/layer.py
@@ -121,7 +121,7 @@ def rasterlayer_reproject(source):
         return layer.srdd.reproject(target_crs, source, resample_method, partition_strategy)
 
     def tiledrasterlayer_reproject(source):
-        return layer.srdd.reproject(target_crs, source, resample_method)
+        return layer.srdd.reproject(target_crs, source, resample_method, partition_strategy)
 
     if isinstance(layer, RasterLayer):
         reproject_method = rasterlayer_reproject