The spark driver has stopped unexpectedly and is restarting

I'm using pyspark on databricks and trying to iterate over a small dataframe (50 rows) using a while loop until the dataframe count is 0 and do some basic manipulations inside the loop (filter, sort, get first row and update some columns) and I get the following error after 15 iterations or so

The spark driver has stopped unexpectedly and is restarting.

After research I found out it's a memory problem.

I'm not using toPandas() or collect(), I'm not using many objects (only 3 dataframes inside the loop and I update them in each iteration), I run the notebook while nothing else is running on the cluster, I tried to increase the driver's memory but I still get the same error. I also tried using gc.collect() and spark.catalog.clearCache() but nothing worked.

The code is something like this:

# df1 and df2 read from the database

columns = ['col1', 'col2', 'col3']

df3 = df1.filter((df1.col2 > 0) & (df1.col3 >= 0))

while df3.count() > 0:

  data = []
  new_data = ['', 0, '']

  df3 = df3.orderBy(df3.col1.asc(), df3.col2.asc())

  df_first = df3.first()

  a = df_first['col1']
  b = df_first['col2']
  c = df_first['col3']

  new_data[0] = a
  new_data[1] = b
  new_data[2] = c

  data.append(tuple(new_data))

  new_df = spark.createDataFrame(data, columns)
  df2 = df2.union(new_df)

  df1 = (df1.withColumn('col2', 
                              when(df1['col1'] == a, df1['col2'] + b)
                              .otherwise(df1['col2'])
                       )
        )

  df3 = (df3.withColumn('col2', 
                              when(df3['col1'] == a, df3['col2'] + b)
                              .otherwise(df3['col2'])
                       )
        )

  df1 = df1.withColumn('col3', df1['col2'] - df1['col4'] - df1['col5'])
  df3 = df3.withColumn('col3', df3['col2'] - df3['col4'] - df3['col5'])

  df3 = df3.filter((df3.col2 > 0) & (df3.col3 >= 0))

Any ideas on what could be causing this issue and how to solve it?