Storing JSON output from multiple individual API calls into Pyspark Dataframe - json

Question

Pyspark newbie here, currently working in Azure Databricks, using a single cluster with 8 cores.
I have a Pyspark dataframe (call this 'df'), where each entry is a property (house/apartment), and has columns for longitude and latitude. The entire dataframe has about 6 million rows.
longitude
latitude
35
97
36
95
...
...
I also have access to an API which, given longitude and latitude, outputs a JSON object containing information about that property (square footage, footprint, roof condition, etc). The output looks something like this ("structures" has a [] around it for some reason):
{"request_id":12345,"structures":[{"footPrint":{"area":275,"wkt":"-83.8 35.9"},"roof_condition":{"value":3.0}],"property_features":{"parcel_wkt":"-83.7 35.9", "requested_wkt":"-83.6 35.9"}}
My question is - How can I most efficiently construct a Pyspark dataframe like the following?
longitude
latitude
request_id
area
wkt
roof_condition
parcel_wkt
requested_wkt
35
97
12345
275
-83.8 35.9
3.0
-83.7 35.9
-83.6 35.9
36
95
12346
280
-81.5 35.4
5.0
-82.0 36.1
-81.1 35.0
...
...
...
...
...
...
...
...
What I have tried -
Create a function to process a single row of the dataframe, then use a UDF with withColumn. For 100 rows, this took over half an hour.
from pyspark.sql.functions import udf
from pyspark.sql.types import *
def process_row(longitude, latitude):
# Use API call to get attributes associated with each lat/lon coordinate
url = 'https://api-url.com/GetInfo?wkt=POINT({} {})&format=json&token=mytoken'.format(longitude, latitude)
data = get_api_info(url) # grabs the output of the API call as a json object
return (data['request_id'], \
data['structures'][0]['footPrint']['area'], \
data['structures'][0]['footprint']['wkt'], \
data['structures'][0]['roof_condition']['value'], \
data['property_features']['parcel_wkt'], \
data['property_features']['requested_wkt'])
schema = StructType([
StructField("request_id", StringType(), True),
StructField("area", FloatType(), True),
StructField("wkt", StringType(), True),
StructField("roof_condition", FloatType(), True),
StructField("parcel_wkt", StringType(), True),
StructField("requested_wkt", StringType(), True)
])
process_row_udf = udf(process_row, schema)
df_ret = df.withColumn('output', process_row_udf('longitude', 'latitude'))
display(df_ret)
Use withColumn to create a single new column 'json_data', which consists of the entire json object as a string. My idea was to then use pyspark.sql.functions.from_json to add new columns from the now existing 'json_data' column to create the request_id, area, wkt, roof_condition, parcel_wkt, and requested_wkt columns. But the first half of this idea was already taking too long (5 minutes for only 30 rows), so I gave up on the rest :(
def process_row2(longitude, latitude):
# Use API call to get attributes associated with each lat/lon coordinate
url = 'https://api-url.com/GetInfo?wkt=POINT({} {})&format=json&token=mytoken'.format(longitude, latitude)
data = get_api_info(url) # grabs the output of the API call as a json object
return str(data)
process_row2_udf = udf(process_row2, StringType())
df_ret2 = df.withColumn('json_data', process_row2_udf('longitude', 'latitude'))
display(df_ret2)
Ideally, I would like to process several rows per second, maybe 3-5, on average. Is there a way to achieve this? I'm confused as to why this is so slow, since I'm pretty sure running a sequential for loop in Pandas is fast than this. I have also checked that the API call itself runs in a fraction of a second, so that's not the bottleneck here. And since I have over 6 million rows to process, I really need this to complete in a reasonable amount of time. Thank you in advance for any help!