IF clause in ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...)

Question

Suppose I have a table like the following:

+----+------+---------+
| id | time | message |
+----+------+---------+
| 1  |  10  |    x    |
| 2  |  12  |    y    |
| 1  |  13  |    z    |
| 2  |  14  |    x    |
| 1  |  15  |    y    |
+----+------+---------+

I want to write a query that returns the most updated message per id. Here is my query:

WITH tmp AS
(
    SELECT *, ROW_NUMBER() OVER (PARTITION BY id ORDER BY time DESC) as rn
    FROM ##TempTable
)
SELECT *
FROM tmp
WHERE rn = 1

which returns:

+----+------+---------+----+
| id | time | message | rn |
+----+------+---------+----+
| 1  |  15  |    y    | 1  |
| 2  |  14  |    x    | 1  |
+----+------+---------+----+

I want to add a condition that, for a specific id, if I see message "z" then just keep that row no matter what the time is, but if "z" is not among the messages keep the most updated row for that id. So, the desired output is like:

+----+------+---------+----+
| id | time | message | rn |
+----+------+---------+----+
| 1  |  13  |    z    | ?  |
| 2  |  14  |    x    | 1  |
+----+------+---------+----+

Any idea how I can modify the query?

Answer 1

here is one way:

select * from (
    select *, row_number() over (partition by id order by case when message = 'z' then 1 else 0 end desc, time desc ) rn
    from data ) t 
where rn = 1

db<>fiddle here