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

由于数据集非常大（~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.frequency 的 dtype 是 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（如x）和frequency（如y）和运行 array_repeat(x, int(y)) 在每个组合上创建 y 重复的数组 x.
展平数组的数组
将 StringType 的一维数组与 space

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

Pyspark: creating string with substring and frequency vector

pyspark

pyspark-sql

pyspark-dataframes