Keep null values as null in PySpark dense_rank()

user_6029749 0
Viewed 863

I'm trying to sort within partitions based on a column that contains nulls.

This DataFrame can be created as follows:

df = spark.createDataFrame( 
                      [("2019-01-01", 5),
                       ("2019-01-01", 20),
                       ("2019-01-01", 10),
                       ("2019-01-03", 1),
                       ("2019-01-03", None),
                       ("2019-01-03", None),
                       ("2019-01-03", 2),
                       ("2019-01-05", 10)],
                       ("Date", "value")
)

Which looks like:

+----------+-----+
|      Date|value|
+----------+-----+
|2019-01-01|    5|
|2019-01-01|   20|
|2019-01-01|   10|
|2019-01-03|    1|
|2019-01-03| null|
|2019-01-03| null|
|2019-01-03|    2|
|2019-01-05|   10|
+----------+-----+

Now I would like to rank each value from 1 to N per date. Which I attempt to do with the following code:

w = Window.partitionBy("date").orderBy("value")
df = df.withColumn("rank", F.dense_rank().over(w))

However, this apparently ranks all null values as 1, regardless of how many null values there are in a column:

+----------+-----+----+
|      Date|value|rank|
+----------+-----+----+
|2019-01-03| null|   1|
|2019-01-03| null|   1|
|2019-01-03|    1|   2|
|2019-01-03|    2|   3|
|2019-01-05|   10|   1|
|2019-01-01|    5|   1|
|2019-01-01|   10|   2|
|2019-01-01|   20|   3|
+----------+-----+----+

How can I change the behavior of dense_rank() such that it keeps null values at null when ranking, instead of giving these values an arbitrary rank of 1?

pyspark python
1 Answers

All you need is an isNull() conditional on both partitioning and the actual ranking. Unfortunately it's not possible to modify dense_rank() directly in pyspark. Code would look like:

from pyspark.sql import Window

common_condition = f.col("value").isNull()

w = Window.partitionBy(
    f.col('Date'),
    f.when(common_condition, 1).otherwise(0)
).orderBy(f.col("value"))

df = df.withColumn(
    "rank", 
    f.when(common_condition, f.lit(None)).otherwise(f.dense_rank().over(w))  
)
df.show()

+----------+-----+----+
|      Date|value|rank|
+----------+-----+----+
|2019-01-03|    1|   1|
|2019-01-03|    2|   2|
|2019-01-03| null|null|
|2019-01-03| null|null|
|2019-01-05|   10|   1|
|2019-01-01|    5|   1|
|2019-01-01|   10|   2|
|2019-01-01|   20|   3|
+----------+-----+----+

Alternatively, you could play around with pyspark.sql.functions.desc_nulls_last() and subsequently filtering out where f.col("value").isNull() and your rank column f.col("rank") is equal to the partitioned f.max() but that is perhaps even clunkier and harder to read.

user_93655760
Related