向 PySpark 中的每一行添加不同的列計數

Question

我需要為 PySpark dataframe 中的每一行添加不同的列計數。

例子：如果原來的dataframe是這樣的：

+----+----+
|col1|col2|
+----+----+
|abc |   1|
|xyz |   1|
|dgc |   2|
|ydh |   3|
|ujd |   1|
|ujx |   3|
+----+----+

然后我想要這樣的東西：

+----+----+----+
|col1|col2|col3|
+----+----+----+
|abc |   1|   3|
|xyz |   1|   3|
|dgc |   2|   3|
|ydh |   3|   3|
|ujd |   1|   3|
|ujx |   3|   3|
+----+----+----+

我試過df.withColumn('total_count', f.countDistinct('col2'))但它給出了錯誤。

Answer 1

您可以計算列中的不同元素並使用以下值創建新列：

distincts = df.dropDuplicates(["col2"]).count()
df = df.withColumn("col3", f.lit(distincts))

Answer 2

交叉連接到不同的計數如下：

df2 = df.crossJoin(df.select(F.countDistinct('col2').alias('col3')))

df2.show()
+----+----+----+
|col1|col2|col3|
+----+----+----+
| abc|   1|   3|
| xyz|   1|   3|
| dgc|   2|   3|
| ydh|   3|   3|
| ujd|   1|   3|
| ujx|   3|   3|
+----+----+----+

Answer 3

您可以使用Window ， collect_set和size ：

from pyspark.sql import functions as F, Window

df = spark.createDataFrame([("abc", 1), ("xyz", 1), ("dgc", 2), ("ydh", 3), ("ujd", 1), ("ujx", 3)], ['col1', 'col2'])

window = Window.orderBy("col2").rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

df.withColumn("col3", F.size(F.collect_set(F.col("col2")).over(window))).show()

+----+----+----+
|col1|col2|col3|
+----+----+----+
| abc|   1|   3|
| xyz|   1|   3|
| dgc|   2|   3|
| ydh|   3|   3|
| ujd|   1|   3|
| ujx|   3|   3|
+----+----+----+