Difference in dense rank and row number in spark

I'm showing @Daniel's answer in Python and I'm adding a comparison with count('*') that can be used if you want to get top-n at most rows per group.

from pyspark.sql.session import SparkSession
from pyspark.sql import Window
from pyspark.sql import functions as F

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame([
    ['a', 10], ['a', 20], ['a', 30],
    ['a', 40], ['a', 40], ['a', 40], ['a', 40],
    ['a', 50], ['a', 50], ['a', 60]], ['part_col', 'order_col'])

window = Window.partitionBy("part_col").orderBy("order_col")

df = (df
  .withColumn("rank", F.rank().over(window))
  .withColumn("dense_rank", F.dense_rank().over(window))
  .withColumn("row_number", F.row_number().over(window))
  .withColumn("count", F.count('*').over(window))
)

df.show()

+--------+---------+----+----------+----------+-----+
|part_col|order_col|rank|dense_rank|row_number|count|
+--------+---------+----+----------+----------+-----+
|       a|       10|   1|         1|         1|    1|
|       a|       20|   2|         2|         2|    2|
|       a|       30|   3|         3|         3|    3|
|       a|       40|   4|         4|         4|    7|
|       a|       40|   4|         4|         5|    7|
|       a|       40|   4|         4|         6|    7|
|       a|       40|   4|         4|         7|    7|
|       a|       50|   8|         5|         8|    9|
|       a|       50|   8|         5|         9|    9|
|       a|       60|  10|         6|        10|   10|
+--------+---------+----+----------+----------+-----+

For example if you want to take at most 4 without randomly picking one of the 4 "40" of the sorting column:

df.where("count <= 4").show()

+--------+---------+----+----------+----------+-----+
|part_col|order_col|rank|dense_rank|row_number|count|
+--------+---------+----+----------+----------+-----+
|       a|       10|   1|         1|         1|    1|
|       a|       20|   2|         2|         2|    2|
|       a|       30|   3|         3|         3|    3|
+--------+---------+----+----------+----------+-----+

In summary, if you filter <= n those columns you will get:

• rank at least n rows
• dense_rank at least n different order_col values
• row_number exactly n rows
• count at most n rows