How can I find the country using latitude & longitude information using Pyspark & shapely

user_9923957 0
Viewed 672

I have two table, one is presenting all countries and their respecting polygons & multipolygons. Another one is presenting latitude and longitude Points() (dist_base).

My goal is here to "join" my point data to their respecting countries. I'm on Pyspark and I can't use geopandas, I'm then using shapely but the way I'm implementing it is really slow.

Do you have any advice using shapely? (the idea would be to check in which polygoin my point is located).

What I tried is to check polygon after polygon (stored in geodata_countries_geo), if the point is inside, if yes, retrieve the ISO3 code associated (stored in a dictionnary geodata_countries_geo_dict)

import shapely.wkt as wkt

def check_nearest_country(lat, lon):
    point = Point(lon, lat)
    for geometry in geodata_countries_geo:
        polygon = wkt.loads(geometry[0])
        if polygon.contains(point):
            return str(geodata_countries_geo_dict.get(str(geometry[0])))
# register udf
check_nearest_country_udf = f.udf(lambda x,y: str(check_nearest_country(x,y)), StringType())
    
# apply
dist_base = dist_base.withColumn('Country_Base_ISO3', check_nearest_country_udf(f.col("lat"),f.col("lon")))
shapely pyspark geospatial python
1 Answers

There is a way to use geopandas in conjunction with Pandas UDFs, which speed up the whole process of spatial join a lot. This presumes that you can collect, through a .toPandas(), your countries dataframe as I think it is not very large - there are not so many countries in the world.

The process is as follows - comments are before each code snippet:

import pandas as pd
import geopandas as gpd
from shapely import wkt
from shapely.geometry import Point
from pyspark.sql.functions import pandas_udf, PandasUDFType



# create sample pyspark.DataFrame of points
df_points = spark.createDataFrame([
  [1.23, 4.56],
  [-1.23, -4.56],
  [0.0, 0.0]
], ['lon', 'lat'])
df_points.show()
# +-----+-----+
# |  lon|  lat|
# +-----+-----+
# | 1.23| 4.56|
# |-1.23|-4.56|
# |  0.0|  0.0|
# +-----+-----+



# create sample geopandas.GeoDataFrame of countries
df_countries = pd.DataFrame({
  'country': ['ITA', 'UK', 'JPN'],
  'polygon': ['POLYGON((1 1,5 1,5 5,1 5,1 1),(2 2, 3 2, 3 3, 2 3,2 2))', 
              'POLYGON((10 10,50 10,50 50,10 50,10 10),(20 20, 30 20, 30 30, 20 30,20 20))',
              'POLYGON((-1 -1,-5 -1,-5 -5,-1 -5,-1 -1),(-2 -2, -3 -2, -3 -3, -2 -3,-2 -2))']
})
geometry = df_countries['polygon'].apply(wkt.loads)
df_countries = df_countries.drop(columns=['polygon'])
gdf_countries = gpd.GeoDataFrame(df_countries, crs="epsg:4326", geometry=geometry)
print(gdf_countries)
#   country                                           geometry
# 0     ITA  POLYGON ((1.00000 1.00000, 5.00000 1.00000, 5....
# 1      UK  POLYGON ((10.00000 10.00000, 50.00000 10.00000...
# 2     JPN  POLYGON ((-1.00000 -1.00000, -5.00000 -1.00000...



# define Pandas UDF
@pandas_udf('string', PandasUDFType.SCALAR)
def spatial_join_udf(lat: pd.Series, lon: pd.Series) -> pd.Series:
    point_var = [Point(xy) for xy in zip(lon, lat)]
    gdf_points = gpd.GeoDataFrame(pd.DataFrame({'lat': lat, 'lon': lon}), crs='epsg:4326', geometry=point_var)
    gdf_joined = gpd.sjoin(gdf_points, gdf_countries, how='left')
    return gdf_joined['country']



# perform spatial join between points and countries
df_points \
  .withColumn('country', spatial_join_udf(df_points['lat'], df_points['lon'])) \
  .show()
# +-----+-----+-------+
# |  lon|  lat|country|
# +-----+-----+-------+
# | 1.23| 4.56|    ITA|
# |-1.23|-4.56|    JPN|
# |  0.0|  0.0|   null|
# +-----+-----+-------+

If you want more information, here's the docs for the function pandas_udf.

user_74654620
Related