Pyspark:使用子字符串和频率向量创建字符串

Pyspark: creating string with substring and frequency vector

我想从一个子字符串列表和一个相应的频率列表创建一个字符串。例如。我的 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'                  |
+--------------------------------+

由于数据集非常大(~22Mio 行),我想尽可能避免 for 循环。 有什么优雅的方法可以做到这一点吗?

非常感谢!

编辑: 我目前的做法:

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 方式。我不知道如何做到这一点。 我还发现 df_in.frequencydtypearray<decimal(4.0)>。所以我试图首先将这些值转换为 int 或在运行时将它们转换为 int

检查以下是否适合您:

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(如x)和frequency(如y)和运行 array_repeat(x, int(y)) 在每个组合上创建 y 重复的数组 x.
  • 展平数组的数组
  • 将 StringType 的一维数组与 space
  • 连接起来