如何在Spark SQL中压缩两个数组列

Question

我有一个熊猫数据框。我尝试将包含字符串值的两列首先连接到列表中，然后使用zip将列表中的每个元素都用'_'连接。我的数据集如下：

df['column_1']: 'abc, def, ghi'
df['column_2']: '1.0, 2.0, 3.0'

我想将这两列连接到第三列，如下所示，用于数据框的每一行。

df['column_3']: [abc_1.0, def_2.0, ghi_3.0]

我已经使用下面的代码在python中成功完成了此操作，但该数据框非常大，并且需要花费很长时间才能为整个数据框运行它。我想在PySpark中做同样的事情以提高效率。我已经成功读取了spark数据框中的数据，但是我很难确定如何使用PySpark等效函数复制Pandas函数。如何在PySpark中获得想要的结果？

df['column_3'] = df['column_2']
for index, row in df.iterrows():
  while index < 3:
    if isinstance(row['column_1'], str):      
      row['column_1'] = list(row['column_1'].split(','))
      row['column_2'] = list(row['column_2'].split(','))
      row['column_3'] = ['_'.join(map(str, i)) for i in zip(list(row['column_1']), list(row['column_2']))]

我已使用以下代码将两列转换为PySpark中的数组

from pyspark.sql.types import ArrayType, IntegerType, StringType
from pyspark.sql.functions import col, split

crash.withColumn("column_1",
    split(col("column_1"), ",\s*").cast(ArrayType(StringType())).alias("column_1")
)
crash.withColumn("column_2",
    split(col("column_2"), ",\s*").cast(ArrayType(StringType())).alias("column_2")
)

现在，我需要使用'_'在两列中压缩数组的每个元素。我该如何使用zip？任何帮助表示赞赏。

Answer 1

与Python等效的Spark SQL为pyspark.sql.functions.arrays_zip：

pyspark.sql.functions.arrays_zip(*cols)

集合函数：返回结构的合并数组，其中第N个结构包含输入数组的所有第N个值。

因此，如果您已经有两个数组：

from pyspark.sql.functions import split

df = (spark
    .createDataFrame([('abc, def, ghi', '1.0, 2.0, 3.0')])
    .toDF("column_1", "column_2")
    .withColumn("column_1", split("column_1", "\s*,\s*"))
    .withColumn("column_2", split("column_2", "\s*,\s*")))

您可以将其应用于结果

from pyspark.sql.functions import arrays_zip

df_zipped = df.withColumn(
  "zipped", arrays_zip("column_1", "column_2")
)

df_zipped.select("zipped").show(truncate=False)

from pyspark.sql.functions import split

df = (spark
    .createDataFrame([('abc, def, ghi', '1.0, 2.0, 3.0')])
    .toDF("column_1", "column_2")
    .withColumn("column_1", split("column_1", "\s*,\s*"))
    .withColumn("column_2", split("column_2", "\s*,\s*")))

我们可以将结果合并transform（？如何使用变换高阶函数，类型错误：列不是可迭代-如何遍历数组类型（）？ ）：

df_zipped_concat = df_zipped.withColumn(
    "zipped_concat",
     expr("transform(zipped, x -> concat_ws('_', x.column_1, x.column_2))")
) 

df_zipped_concat.select("zipped_concat").show(truncate=False)

from pyspark.sql.functions import arrays_zip

df_zipped = df.withColumn(
  "zipped", arrays_zip("column_1", "column_2")
)

df_zipped.select("zipped").show(truncate=False)

注意事项：

高阶功能transform，并arrays_zip已在Apache中星火2.4引入。

Answer 2

zip_with对于 Spark 2.4+，可以仅使用同时压缩连接的函数来完成此操作：

df.withColumn("column_3", expr("zip_with(column_1, column_2, (x, y) -> concat(x, '_', y))")) 

高阶函数需要使用 lambda 函数按元素合并 2 个数组(x, y) -> concat(x, '_', y)。