Pyspark: создание строки с подстрокой и вектором частоты

Question

Я хочу создать строку из списка подстрок и соответствующего списка частот. Например, мой df_in выглядит следующим образом:

+-------------------------+-----------+
|         substr          | frequency |
+-------------------------+-----------+
| ['ham', 'spam', 'eggs'] | [1, 2, 3] |
| ['foo', 'bar']          | [2, 1]    |
+-------------------------+-----------+

И я хочу, чтобы мой df_out выглядел следующим образом:

+--------------------------------+
|             output             |
+--------------------------------+
| 'ham spam spam eggs eggs eggs' |
| 'foo foo bar'                  |
+--------------------------------+

Поскольку набор данных очень большой (~ 22 миллионов строк)Я хотел бы избежать петель для всех возможных. Есть ли какой-нибудь элегантный способ добиться этого?

Большое спасибо!

Редактировать: Мой текущий подход:

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))

Проблема:

Я прочиталчто для ускорения вычислений UDF должны быть переписаны в pyspark. Я не знаю, как это можно сделать, хотя. Также я узнал, что dtype из df_in.frequency - это array<decimal(4.0)>. Поэтому я пытаюсь либо сначала преобразовать эти значения в int, либо привести их к int во время выполнения.

Answer 1

Проверьте, работает ли для вас следующее:

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                 |
+-----------------+---------+----------------------------+

Ниже описано, как это работает:

• используйте zip_with для итерации двух массивов substr (как * 1009)*) и frequency (как y) бок о бок и запустите array_repeat(x, int(y)) для каждой комбинации, чтобы создать массив y повторов x.
• , чтобы сгладить массив массивов
• объединить одномерный массив StringType с пробелом