Skip to content

Commit

Permalink
[SEDONA-712] add toSpatialRdd method for RDD[InternalRow] (#1816)
Browse files Browse the repository at this point in the history
  • Loading branch information
james-willis authored Feb 14, 2025
1 parent 41a6bc0 commit 6a82808
Showing 1 changed file with 18 additions and 5 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -84,17 +84,30 @@ object StructuredAdapter {
}

/**
* Convert DataFrame to SpatialRDD. It puts InternalRow as user data of Geometry. It allows only
* one geometry column.
* Convert DataFrame to SpatialRDD. It puts InternalRow as user data of Geometry.
*
* @param dataFrame
* @param geometryFieldName
*/
def toSpatialRdd(dataFrame: DataFrame, geometryFieldName: String): SpatialRDD[Geometry] = {
def toSpatialRdd(dataFrame: DataFrame, geometryFieldName: String): SpatialRDD[Geometry] =
toSpatialRdd(dataFrame.queryExecution.toRdd, dataFrame.schema, geometryFieldName)

/**
* Convert RDD[InternalRow] to SpatialRDD. It puts InternalRow as user data of Geometry.
*
* @param rdd
* @param schema
* @param geometryFieldName
* @return
*/
def toSpatialRdd(
rdd: RDD[InternalRow],
schema: StructType,
geometryFieldName: String): SpatialRDD[Geometry] = {
val spatialRDD = new SpatialRDD[Geometry]
spatialRDD.schema = dataFrame.schema
spatialRDD.schema = schema
val ordinal = spatialRDD.schema.fieldIndex(geometryFieldName)
spatialRDD.rawSpatialRDD = dataFrame.queryExecution.toRdd
spatialRDD.rawSpatialRDD = rdd
.map(row => {
val geom = GeometrySerializer.deserialize(row.getBinary(ordinal))
geom.setUserData(row.copy())
Expand Down

0 comments on commit 6a82808

Please sign in to comment.