Question

我有一个PySpark DataFrame，结构由

给出

[('u1', 1, [1 ,2, 3]), ('u1', 4, [1, 2, 3])].toDF('user', 'item', 'fav_items')

我需要添加一个1或0的列，具体取决于'item'是否在'fav_items'中。

所以我想要

[('u1', 1, [1 ,2, 3], 1), ('u1', 4, [1, 2, 3], 0)]

我如何查找第二列到第三列来决定值以及如何添加它？

Answer 1

以下代码执行请求的任务。定义了一个用户定义的函数，它接收两列['error']作为参数。因此，对于每一行，搜索项目是否在项目列表中。如果找到该项，则返回1，否则返回0.

DataFrame

结果如下：

# Imports
from pyspark.sql.types import IntegerType
from pyspark.sql.functions import udf
# First we create a RDD in order to create a dataFrame:
rdd = sc.parallelize([('u1', 1, [1 ,2, 3]), ('u1', 4, [1, 2, 3])])
df = rdd.toDF(['user', 'item', 'fav_items'])
# Print dataFrame
df.show()

# We make an user define function that receives two columns and do operation
function = udf(lambda item, items: 1 if item in items else 0, IntegerType())

df.select('user', 'item', 'fav_items', function(col('item'), col('fav_items')).alias('result')).show()

Answer 2

只是为了有趣的非UDF解决方案：

from pyspark.sql.functions import col, first, explode, max as max_

result = (
    # Here we take exploded rows and for each row check if there
    # is a match. We cast to integer (false -> 0, true -> 1)
    # and take max (1 if there is any match)
    max_((col("fav_item") == col("item")).cast("integer"))
).alias("result")


(df.repartition("user", "item") 
  # Explode array so we compare item and fav_item
  .withColumn("fav_item", explode("fav_items")) 
  .groupBy("user", "item")
  # Aggregate
  # we add result and retain fav_items
  .agg(result, first("fav_items").alias("fav_items")))

所以它只是：

展开fav_array：

## +----+----+---------+--------+
## |user|item|fav_items|fav_item|
## +----+----+---------+--------+
## |  u1|   1|[1, 2, 3]|       1|
## |  u1|   1|[1, 2, 3]|       2|
## |  u1|   1|[1, 2, 3]|       3|
## |  u1|   4|[1, 2, 3]|       1|
## |  u1|   4|[1, 2, 3]|       2|
## |  u1|   4|[1, 2, 3]|       3|
## +----+----+---------+--------+

检查fav_item = item（_1是(col("fav_item") == col("item")).cast("integer")表达式的结果）：

## +----+----+---------+--------+---+
## |user|item|fav_items|fav_item| _1|
## +----+----+---------+--------+---+
## |  u1|   1|[1, 2, 3]|       1|  1|
## |  u1|   1|[1, 2, 3]|       2|  0|
## |  u1|   1|[1, 2, 3]|       3|  0|
## |  u1|   4|[1, 2, 3]|       1|  0|
## |  u1|   4|[1, 2, 3]|       2|  0|
## |  u1|   4|[1, 2, 3]|       3|  0|
## +----+----+---------+--------+---+

并将其回滚，保留user和item作为组列，任意fav_items（都相同）和临时列_1的最大值（ 0或1）。

我会选择使用UDF。

根据列值是否在另一列中，将列添加到PySpark DataFrame

2 个答案: