What is the optimal way to merge many Dataframes in PySpark?

Viewed 105

I have N (100-1K) of huge Dataframes (100GB - 1T each) and need to process each of them and get a small dataframe (couple of MB) and union the processed dataframes. Is there a more optimal way of doing it that this:

df_merged = new_empty_df

for i in range(N):
    df_raw = read(i)
    df_processed = process(df_raw)
    df_merged = df_merged.union(df_processed) # should .persist() be used here?

# output processed
df_merged.write(somewhere)

Also, should .persist() be used in the 3rd line of the loop to speed things up?

Edit 1: Here' more information about the data and what's happening inside the process. Data is tsv/parquet containing GPS locations:

Location
POINT (30 10)
POINT (33.44 -22.4)
POINT (33.11 -21.7)

The process first calculates AreaId for each Location point. It's applied to every row.

Location AreaId
POINT (30 10) 2
POINT (33.44 -22.4) 4
POINT (33.11 -21.7) 4

Then, I need want to calculate the number of points inside each AreaId, resulting in the following:

AreaId Count
2 1
4 2

The code I used for those 2 things is:

df_area_id= df_raw.withColumn('AreaId',GetTileForPointUDF(col('Location')))
        
df_processed = df_area_id.groupby('AreaId').count()
1 Answers

There is no need in persist as you are using your dataframes only once. Your query is already optimal. For example, a similar query gives the following plan.

def process(df):
 return  df.withColumn("id", F.col("id") + 1)

df_merged = spark.range(0)
N = 3

for i in range(N):
    df_raw = spark.range(i+1)
    df_processed = process(df_raw)
    df_merged = df_merged.union(df_processed)

 == Physical Plan ==
     Union
     :- *(1) Range (0, 0, step=1, splits=400)     // new_empty_df
     :- *(2) Project [(id#132L + 1) AS id#134L]   // process raw 1
     :  +- *(2) Range (0, 1, step=1, splits=400)  // read raw 1
     :- *(3) Project [(id#137L + 1) AS id#139L]   // process raw 2
     :  +- *(3) Range (0, 2, step=1, splits=400)  // read raw 2
     +- *(4) Project [(id#142L + 1) AS id#144L]   // process raw 3
        +- *(4) Range (0, 3, step=1, splits=400)  // read raw 3
Related