Fixes #67

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

Author: metasim

core/src/main/scala/astraea/spark/rasterframes/extensions/ContextRDDMethods.scala

@@ -39,8 +39,9 @@ abstract class SpatialContextRDDMethods[T <: CellGrid](implicit spark: SparkSess
   def toRF(implicit converter: PairRDDConverter[SpatialKey, T]): RasterFrame = toRF(TILE_COLUMN.columnName)
 
   def toRF(tileColumnName: String)(implicit converter: PairRDDConverter[SpatialKey, T]): RasterFrame = {
-    val df = self.toDataFrame
-    df.setSpatialColumnRole(SPATIAL_KEY_COLUMN, self.metadata)
+    val df = self.toDataFrame.setSpatialColumnRole(SPATIAL_KEY_COLUMN, self.metadata)
+    val defName = TILE_COLUMN.columnName
+    df.mapWhen(_ ⇒ tileColumnName != defName, _.withColumnRenamed(defName, tileColumnName))
       .certify
   }
 }
@@ -56,9 +57,11 @@ abstract class SpatioTemporalContextRDDMethods[T <: CellGrid](
   def toRF(implicit converter: PairRDDConverter[SpaceTimeKey, T]): RasterFrame = toRF(TILE_COLUMN.columnName)
 
   def toRF(tileColumnName: String)(implicit converter: PairRDDConverter[SpaceTimeKey, T]): RasterFrame = {
-    self.toDataFrame
+    val df = self.toDataFrame
       .setSpatialColumnRole(SPATIAL_KEY_COLUMN, self.metadata)
       .setTemporalColumnRole(TEMPORAL_KEY_COLUMN)
+    val defName = TILE_COLUMN.columnName
+    df.mapWhen(_ ⇒ tileColumnName != defName, _.withColumnRenamed(defName, tileColumnName))
       .certify
   }
 }

core/src/test/scala/astraea/spark/rasterframes/RasterFrameSpec.scala

@@ -74,6 +74,10 @@ class RasterFrameSpec extends TestEnvironment with MetadataKeys
       )
 
       assert(rf.count() === 4)
+
+      val cols = tileLayerRDD.toRF("foo").columns
+      assert(!cols.contains("tile"))
+      assert(cols.contains("foo"))
     }
 
     it("should implicitly convert from spatiotemporal layer type") {
@@ -93,6 +97,9 @@ class RasterFrameSpec extends TestEnvironment with MetadataKeys
           println(rf.schema.prettyJson)
           throw ex
       }
+      val cols = tileLayerRDD.toRF("foo").columns
+       assert(!cols.contains("tile"))
+       assert(cols.contains("foo"))
     }
 
     it("should implicitly convert layer of TileFeature") {