将 PySpark DataFrame ArrayType 字段组合成单个 ArrayType 字段

Question

我有一个带有 2 个 ArrayType 字段的 PySpark DataFrame：

 >>>df
DataFrame[id: string, tokens: array<string>, bigrams: array<string>]
>>>df.take(1)
[Row(id='ID1', tokens=['one', 'two', 'two'], bigrams=['one two', 'two two'])]

我想将它们组合成一个 ArrayType 字段：

 >>>df2
DataFrame[id: string, tokens_bigrams: array<string>]
>>>df2.take(1)
[Row(id='ID1', tokens_bigrams=['one', 'two', 'two', 'one two', 'two two'])]

适用于字符串的语法在这里似乎不起作用：

 df2 = df.withColumn('tokens_bigrams', df.tokens + df.bigrams)

谢谢！

原文由 zemekeneng 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

火花 >= 2.4

您可以使用 concat 函数 ( SPARK-23736 )：

 from pyspark.sql.functions import col, concat

df.select(concat(col("tokens"), col("tokens_bigrams"))).show(truncate=False)

# +---------------------------------+
# |concat(tokens, tokens_bigrams)   |
# +---------------------------------+
# |[one, two, two, one two, two two]|
# |null                             |
# +---------------------------------+

要在其中一个值为 NULL 时保留数据，您可以 coalesce 与 array ：

 from pyspark.sql.functions import array, coalesce

df.select(concat(
    coalesce(col("tokens"), array()),
    coalesce(col("tokens_bigrams"), array())
)).show(truncate = False)

# +--------------------------------------------------------------------+
# |concat(coalesce(tokens, array()), coalesce(tokens_bigrams, array()))|
# +--------------------------------------------------------------------+
# |[one, two, two, one two, two two]                                   |
# |[three]                                                             |
# +--------------------------------------------------------------------+

火花 < 2.4

不幸的是连接 array 列在一般情况下你需要一个 UDF，例如这样的：

 from itertools import chain
from pyspark.sql.functions import col, udf
from pyspark.sql.types import *

def concat(type):
    def concat_(*args):
        return list(chain.from_iterable((arg if arg else [] for arg in args)))
    return udf(concat_, ArrayType(type))

可以用作：

 df = spark.createDataFrame(
    [(["one", "two", "two"], ["one two", "two two"]), (["three"], None)],
    ("tokens", "tokens_bigrams")
)

concat_string_arrays = concat(StringType())
df.select(concat_string_arrays("tokens", "tokens_bigrams")).show(truncate=False)

# +---------------------------------+
# |concat_(tokens, tokens_bigrams)  |
# +---------------------------------+
# |[one, two, two, one two, two two]|
# |[three]                          |
# +---------------------------------+

原文由 zero323 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

在 Spark 2.4.0（Databricks 平台上为 2.3）中，您可以使用 concat 函数在 DataFrame API 中本地执行此操作。在您的示例中，您可以这样做：

 from pyspark.sql.functions import col, concat

df.withColumn('tokens_bigrams', concat(col('tokens'), col('bigrams')))

这 是相关的jira。

原文由 David Vrba 发布，翻译遵循 CC BY-SA 4.0 许可协议