[英]Convert a PySpark Dataframe Column to a Python List depending on the value in another column
[英]Adding column to PySpark DataFrame depending on whether column value is in another column
我有一个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)]
我如何查找第二列到第三列来决定值以及如何添加它?
以下代码执行所请求的任务。 定义了一个用户定义的函数,它接收两列DataFrame
作为参数。 因此,对于每一行,搜索项目是否在项目列表中。 如果找到该项,则返回1,否则返回0。
# 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()
结果如下:
+----+----+---------+
|user|item|fav_items|
+----+----+---------+
| u1| 1|[1, 2, 3]|
| u1| 4|[1, 2, 3]|
+----+----+---------+
+----+----+---------+------+
|user|item|fav_items|result|
+----+----+---------+------+
| u1| 1|[1, 2, 3]| 1|
| u1| 4|[1, 2, 3]| 0|
+----+----+---------+------+
只是为了有趣的非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")
的结果(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。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.