Geospatial Joins in Apache Sedona

Pyspark combined with Apache Sedona is a very powerful tool that not enough data scientists and data engineers are using, especially when it comes to spatial data. This walkthrough will tell you the necessities for getting started in using spatial joins with Apache Sedona in Databricks.

First, sign up for a Community Edition Databricks account When you are logged in create a new notebook and start a cluster with a 12.2 LTS (Scala 2.12, Spark 3.3.2) runtime, then attach it to the notebook.

Click on “Compute” on the left nav-bar and click on your cluster. Click on the “Libraries” tab and click “Install new”.

Select Maven and enter the following in the coordinates section:

org.apache.sedona:sedona-spark-shaded-3.0_2.12:1.6.0

Click “Install”.

Click the install new button again and select “Maven”.

Enter the following into the coordinates section:

org.datasyslab:geotools-wrapper:1.6.0-28.2

Go back to your notebook and you will use the following line to install the Apache Sedona library for Python:

%pip install apache-sedona

Then import the library and initialize the Sedona context:

from pyspark.sql.types import StructType, StructField, FloatType, TimestampType, ShortType, DoubleType, DateType, StringType, LongType
from pyspark.sql import functions as F
from sedona.spark import *

config = SedonaContext.builder() .\
    config('spark.jars.packages',
           'org.apache.sedona:sedona-spark-shaded-3.0_2.12:1.6.0,'
           'org.datasyslab:geotools-wrapper:1.6.0-28.2'). \
    getOrCreate()

sedona = SedonaContext.create(config)

Now you’ll make two sets of data from the example below. The first uses points and has columns id, latitude, and longitude. The second has polygons with columns id and WKT.

points_data = [
    (1, 34.05, -118.25),  # Los Angeles
    (2, 36.16, -115.15),  # Las Vegas
    (3, 40.71, -74.00),   # New York
    (4, 37.77, -122.41)   # San Francisco
]


polygons_data = [
    (1, "POLYGON((-125 30, -125 40, -115 40, -115 30, -125 30))"),  # Covers part of California/Nevada
    (2, "POLYGON((-80 35, -80 45, -70 45, -70 35, -80 35))")        # Covers part of New York area
]

Define the schemas for the points data and polygons.

points_schema = StructType([
    StructField("id", StringType(), True),
    StructField("latitude", DoubleType(), True),
    StructField("longitude", DoubleType(), True)
])

polygons_schema = StructType([
    StructField("id", StringType(), True),
    StructField("wkt", StringType(), True)
])

Next, create DataFrames from the sample data.

points_df = spark.createDataFrame(points_data, schema=points_schema)
polygons_df = spark.createDataFrame(polygons_data, schema=polygons_schema)

Convert points from the first data frame to Geometry objects.

points_df = points_df.withColumn("geometry", F.expr("ST_Point(CAST(longitude AS Decimal(24, 20)), CAST(latitude AS Decimal(24, 20)))"))

Convert the polygons from the second data frame to Geometry objects.

polygons_df = polygons_df.withColumn("geometry", F.expr("ST_GeomFromWKT(wkt)"))

Now we are going to switch from using Pyspark to using SparkSQL to take advantage of SQL’s simple join operation.

Start by registering the DataFrames as temporary views for SQL operations.

points_df.createOrReplaceTempView("points")
polygons_df.createOrReplaceTempView("polygons")

Perform geospatial join using ST_Intersects, which joins the rows between data frames based on if their geometries share any area at all.

result_df = spark.sql("""
    SELECT p.id AS point_id, g.id AS polygon_id
    FROM points p, polygons g
    WHERE ST_Intersects(p.geometry, g.geometry)
""")

ST_Intersects is just one of many useful Apache Sedona spatial functions.

Thanks for reading this brief guide on getting started with spatial joins in Apache Sedona and Databricks.