将PySpark DataFrame ArrayType字段组合到单个ArrayType字段中

时间:2016-05-17 18:48:58

标签: python apache-spark dataframe pyspark apache-spark-sql

我有一个包含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)

谢谢!

3 个答案:

答案 0 :(得分:21)

Spark&gt; = 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]                                                             |
# +--------------------------------------------------------------------+

Spark&lt; 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]                          |
# +---------------------------------+

答案 1 :(得分: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')))

Here是相关的吉拉。

答案 2 :(得分:0)

我使用的是 Spark < 2.4,但上述解决方案对我不起作用,出现错误“函数 concat 的输入应该具有 StringType 或 BinaryType”。 这对我有用:

from pyspark.sql import functions as F

df.select("*",F.array(F.concat_ws(',', col('tokens'), col('bigrams))).\
                            alias('concat_cols'))