Perform regex_replace on pyspark dataframe using multiple dictionaries containing specific key/value pairs without looping

We need to parse some text data in several very large dataframes. Part of this parsing involves utilising a set of dictionaries, each containing a series of key/value pairs that represent words or characters to be replaced, and their replacements. The dictionaries must be used in order, and cannot be combined. This is a requirement. We are encountering several problems with our implementation, which currently uses PySpark regex_replace and loops in the following way (dictionaries and data are not this simple, but the example is representative):

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
import pandas as pd

spark = SparkSession.builder.appName('regex_parser').getOrCreate()

cols = ['input']
data = ['cat', 'toad', 'pig', 'abc', '123', '{.-!']

dict1 = {'cat': 'dog',
         'toad': 'frog',
         'pig': 'hog'}

dict2 = {'a': 'A',
         'b': 'B',
         'c': 'Z',
         '{': '-',
         '.': '-'}


df = spark.createDataFrame(pd.DataFrame(columns=cols, data=data))

df = df.withColumn('dict1', F.col('input'))
current_dict = dict1
for k, v in current_dict.items():
    df1 = df1.withColumn('dict1', F.regex_replace(F.col('dict1'), k, v))

df = df.withColumn('dict2', F.col('dict1'))
current_dict = dict2
for k, v in current_dict.items():
    df1 = df1.withColumn('dict2', F.regex_replace(F.col('dict2'), k, v))

df.show()

This method works in general, but when our full implantation (a total of 9 dictionaries, equalling approx 1000 unique replacement pairs) is run, the execution plan becomes so long that all executors are dropped and the session crashes, even when we include config commands to increase the yarn and pyspark memory overheads, set codegen.wholeStage to false and disable dynamic allocation of executors. Looking at the DAG, we see 1000 nested regex calls.

To fix this, we implemented the following change every few loops:

df = df.withColumn('dict1', F.col('input'))
current_dict = dict1
for k, v in current_dict.items():
    df1 = df1.withColumn('dict1', F.regex_replace(F.col('dict1'), k, v))

df.cache().count()

df = df.withColumn('dict2', F.col('dict1'))
current_dict = dict2
for k, v in current_dict.items():
    df1 = df1.withColumn('dict2', F.regex_replace(F.col('dict2'), k, v))

This forces the plan to execute when it is smaller, and the run successfully completes most of the time. However, executors are still dropped regularly (and picked back up again by the yarn scheduler) if dynamic allocation is enabled, and sometimes (depending on what other processes are being run on our cluster at the time), the job fails entirely. We could set the dynamic allocation to false, but the parser is just the first step of a much larger program, and if executors are dropped later, they will not be returned, causing the job to fail.

The variant of the code with periodic df.cache().count() statements is the current version. It works most of the time, though we are alarmed by the amount of executors that fall over and are picked back up. It also fails with much more regularity on our full dataset, which is many times larger than our test data. In addition, the DAG still reads as 1000 nested regex_replace statements, though the caching insures that not all are executed at once.

We have attempted to use the following to break the execution plan chain instead of df.cache().count():

df.checkpoint()
df.localCheckpoint()
df.write.mode('overwrite').parquet('foobar')

All 3 variants fail due to an issue where the file is still being written out when the next operation is being executed. We cant seem to fix this.

Is there a way of parsing these dictionaries in without using loops? Unfortunately, we cannot use udfs of any description. This is also a requirement (otherwise this becomes very easy with pandas and pyarrow!)