Pyspark：查询自身数据框以创建新列

Question

我在 pyspark 中有以下数据集

Id	Sub
1	Mat
1	Phy
1	Sci
2	Bio
2	Phy
2	Sci

我想创建一个类似于下面的 df

Id	Sub	HaMat
1	Mat	1
1	Phy	1
1	Sci	1
2	Bio	0
2	Phy	0
2	Sci	0

我如何在 pyspark 中执行此操作？

def hasMath(studentID,df):
    return df.filter(col('Id') == studentID & col('sub') = 'Mat' ).count()

df = df.withColumn("hasMath",hasMath(F.col('id'),df1))

但这似乎不起作用。有没有更好的方法来实现这个。

Answer 1

您可以在 window 和 expr（对于 Spark 2.4+）上使用 collect_list 来获取每个 ID 的主题列表并过滤 Mat.

size函数获取数组的计数。

from pyspark.sql import functions as F
from pyspark.sql.functions import col
from pyspark.sql.window import Window

w = Window().partitionBy("Id")

df.withColumn("list", F.collect_list(col("Sub")).over(w))\
  .withColumn("hasMath", F.size(F.expr("filter(list, x -> x == 'Mat')")))\
  .drop("list").show()

输出：

+---+---+-------+
| Id|Sub|hasMath|
+---+---+-------+
|  1|Phy|      1|
|  1|Mat|      1|
|  1|Sci|      1|
|  2|Phy|      0|
|  2|Bio|      0|
|  2|Sci|      0|
+---+---+-------+

Pyspark：查询自身数据框以创建新列

Pyspark : Query self dataframe to create a new column

pyspark