缺少列值时将该列值添加到另一列的数组中

时间:2019-05-30 22:38:44

标签: python pyspark

我有一个数据框,其中的一列包含一个数组列表,该列表包含格式(键,分数)的结构,即下面的值列。 另外,我还有另一个名为item的列,其中包含一个字符串。

root
 |-- value: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- key: string (nullable = true)
 |    |    |-- score: double (nullable = true)
 |-- id: string (nullable = true)
 |-- item: string (nullable = true)

当该项目不为null时,我想搜索arraylist以查看value.key(s)中是否存在该项目。如果不存在,则将(项目,最小)附加到value列中的arraylist中。在此,min是数组列表中分数的最小值。 我需要保留arraylist中元素的顺序,并将新结构附加到列表的末尾。

样本输入:

+-----------------------------------+------+------+
| value                             |  id  | item |
+-----------------------------------+------+------+
|[[e1, 0.4] , [e2, 0.3]]            | 1    | e3   |
|[[e1, 0.4] , [e2, 0.3], [e3, 0.2]] | 2    | e4   |
|[[e1, 0.4] , [e2, 0.3]]            | 3    | e1   |
|[[e1, 0.4] , [e2, 0.3]]            | 4    | null |
+-----------------------------------+------+------+

示例输出:

+------------------------------------------------+------+------+
| value                                          |  id  | item |
+------------------------------------------------+------+------+
|[[e1, 0.4] , [e2, 0.3], [e3, 0.3]]              | 1    | e3   |
|[[e1, 0.4] , [e2, 0.3], [e3, 0.2], [e4, 0.2]]   | 2    | e4   |
|[[e1, 0.4] , [e2, 0.3]]                         | 3    | e1   |
|[[e1, 0.4] , [e2, 0.3]]                         | 4    | null |
+------------------------------------------------+------+------+

如果我想将新的(键,分数)添加到列表中的特定位置,解决方案会有什么不同? (例如开始或中间)

1 个答案:

答案 0 :(得分:0)

您可以定义udf函数来实现它。

import pyspark.sql.functions as F
from pyspark.sql.types import StructType,StructField,ArrayType,StringType,DoubleType

def contains(values,item):
    if not item:
        return values
    keys = [pair['key'] for pair in values]
    if item not in keys:
        scores = [pair['score'] for pair in values]
        values.append({'key':item,'score':min(scores)})
    return values

contains_udf = F.udf(contains,ArrayType(StructType([StructField('key', StringType()),
                                                    StructField('score', DoubleType())])))
df = df.withColumn("value", contains_udf('value','item'))
df.show(truncate = False)

+--------------------------------------------+---+----+
|value                                       |id |item|
+--------------------------------------------+---+----+
|[[e1, 0.4], [e2, 0.3], [e3, 0.3]]           |1  |e3  |
|[[e1, 0.4], [e2, 0.3], [e3, 0.2], [e4, 0.2]]|2  |e4  |
|[[e1, 0.4], [e2, 0.3]]                      |3  |e1  |
|[[e1, 0.4], [e2, 0.3]]                      |4  |null|
+--------------------------------------------+---+----+