Remove the number of duplicates which is more than 4. Only a maximum of 4 duplicates will remain

Viewed 43

I have a DataFrame with duplicates in col text (in reality, the content of this column is a concat of several other columns). Col dubs is the number of duplicates for col text.

My challange:
Remove the number of duplicates which is more than 4. So only a maximum of 4 duplicates will remain. If less or equal to 4 duplicates, keep all duplicates. There are no other criteria.

df = spark.createDataFrame([
  ["aaa","1"],
  ["bbb","2"],
  ["bbb","2"],
  ["ccc","5"],
  ["ccc","5"],
  ["ccc","5"],
  ["ccc","5"],
  ["ccc","5"],
  ["ddd","4"],
  ["ddd","4"],
  ["ddd","4"],
  ["ddd","4"]
]).toDF("text","dubs")

Expected output:

+----+----+
|text|dubs|
+----+----+
| aaa|   1|
| bbb|   2|
| bbb|   2|
| ccc|   5|
| ccc|   5|
| ccc|   5|
| ccc|   5|
| ddd|   4|
| ddd|   4|
| ddd|   4|
| ddd|   4|
+----+----+

I know how to remove duplicates, but this is a bit mor challenging for me and I really have no idea how to handle this. Any help is greatly appreciated!

1 Answers

You can assign row numbers for partitioned text column and remove any row that has row number > 4 against any text column value.

df.withColumn("row_num", row_number().over(Window.partitionBy("text").orderBy("text"))) \
  .filter(col("row_num")< 5).drop("row_num") \
  .orderBy("text").show()



+----+----+
|text|dubs|
+----+----+
| aaa|   1|
| bbb|   2|
| bbb|   2|
| ccc|   5|
| ccc|   5|
| ccc|   5|
| ccc|   5|
| ddd|   4|
| ddd|   4|
| ddd|   4|
| ddd|   4|
+----+----+
Related