How to select different rows base on conditions in Spark sql

I created a rank column base on num column and grouped by name. I'm trying to implement logic to select the colour value in the first row if num is not 0, otherwise use the colour value from the second row or 3rd row, where the num is not 0 For example for Name A, colour yellow should be selected. For Name B, colour green should be selected.
Could someone please help to suggest how to do this? Thank you.

Name.    colour     num.     rank 
 A.       blue       0         1
 A        yellow.    300       2
 B        green      100.      1
 B        brown      500.      2     

You can use first use filter to exclude 0s, followed by a Window partition to assign row_number, and then another filter to filter out values thare are not 1.

Example

df
  .filter(col("num").gt(0))
  .withColumn("row_number", row_number().over(Window.partitionBy("name").orderBy(col("rank"))))
  .filter(col("row_number").equalTo(1))
  .drop("row_number")

Output

+----+------+---+----+
|name|color |num|rank|
+----+------+---+----+
|A   |yellow|300|2   |
|B   |green |100|1   |
+----+------+---+----+

This only works under the assumption that the combination of name and num are unique, if this is not the case, then further adjustments are needed. Also, if you show the code that leads to your mentioned table, this problem could be solved in a more efficient way (probably). Good luck!