Pyspark: creating string with substring and frequency vector - apache-spark-sql

Question

I want to create a string from a list of substrings and a correponding frequency list. E.g. my df_in looks like following:
+-------------------------+-----------+
| substr | frequency |
+-------------------------+-----------+
| ['ham', 'spam', 'eggs'] | [1, 2, 3] |
| ['foo', 'bar'] | [2, 1] |
+-------------------------+-----------+
And I want my df_out to look like this:
+--------------------------------+
| output |
+--------------------------------+
| 'ham spam spam eggs eggs eggs' |
| 'foo foo bar' |
+--------------------------------+
Since the dataset is very large (~22Mio rows), I would like to avoid for-loops whereever possible.
Is there any elegant way to achieve this?
Thanks a lot!
Edit:
My current approach:
import pyspark.sql.functions as F
import pyspark.sql.types as T
def create_text(l_sub, l_freq):
l_str = [(a+' ')*b if isinstance(b, int) else (a+' ') for a, b in zip(l_sub, l_freq)]
return ''.join(l_str)
create_str = F.udf(lambda x, y: create_text(x, y), T.StringType())
df = df.withColumn('output', create_str(df_in.sbustr, df_in.frequency))
Problem:
I read that in order to speed up the compute, UDFs should be re-written into pyspark way. I do not know how that can be done though.
Also I found out the dtype of df_in.frequency is array<decimal(4.0)>. So I am trying to either convert these values to int first or cast them to int in runtime.

Answer 1

Check if the following works for you:
from pyspark.sql.functions import expr
df.withColumn('output', expr('''
array_join(flatten(zip_with(`substr`, `frequency`, (x,y) -> array_repeat(x,int(y)))), ' ')
''')).show(truncate=False)
+-----------------+---------+----------------------------+
|substr |frequency|output |
+-----------------+---------+----------------------------+
|[ham, spam, eggs]|[1, 2, 3]|ham spam spam eggs eggs eggs|
|[foo, bar] |[2, 1] |foo foo bar |
+-----------------+---------+----------------------------+
Below is how it works:
use zip_with to iterate two arrays substr(as x) and frequency(as y) side-by-side and run array_repeat(x, int(y)) on each combination to create an array of y repeats of x.
flatten the array of arrays
join the 1-D array of StringType with space