Translating a SAS Ranking with Tie set to HIGH into PySpark

Viewed 196

I'm trying to replicate the following SAS code in PySpark:

PROC RANK DATA = aud_baskets OUT = aud_baskets_ranks GROUPS=10 TIES=HIGH;
BY customer_id;
VAR expenditure;
RANKS basket_rank;
RUN;

The idea is to rank all expenditures under each customer_id block. The data would look like this:

+-----------+--------------+-----------+
|customer_id|transaction_id|expenditure|
+-----------+--------------+-----------+
|          A|             1|         34|
|          A|             2|         90|
|          B|             1|         89|
|          A|             3|          6|
|          B|             2|          8|
|          B|             3|          7|
|          C|             1|         96|
|          C|             2|          9|
+-----------+--------------+-----------+

In PySpark, I tried this:

spendWindow = Window.partitionBy('customer_id').orderBy(col('expenditure').asc())
aud_baskets = (aud_baskets_ranks.withColumn('basket_rank', ntile(10).over(spendWindow)))

The problem is that PySpark doesn't let the user change the way it will handle Ties, like SAS does (that I know of). I need to set this behavior in PySpark so that values are moved up to the next tier each time one of those edge cases occur, as oppose to dropping them to the rank below.

Or is there a way to custom write this approach?

2 Answers

Use dense_rank it will give same rank in case of ties and next rank will not be skipped ntile function split the group of records in each partition into n parts. In your case which is 10

from pyspark.sql.functions import dense_rank

spendWindow = Window.partitionBy('customer_id').orderBy(col('expenditure').asc())
aud_baskets = aud_baskets_ranks.withColumn('basket_rank',dense_rank.over(spendWindow))

Try The following code. It is generated by an automated tool called SPROCKET. It should take care of ties.

df = (aud_baskets)
for (colToRank,rankedName) in zip(['expenditure'],['basket_rank']): 
    wA = Window.orderBy(asc(colToRank))
    df_w_rank = (df.withColumn('raw_rank', rank().over(wA)))
    ties = df_w_rank.groupBy('raw_rank').count().filter("""count > 1""")
    df_w_rank = (df_w_rank.join(ties,['raw_rank'],'left').withColumn(rankedName,expr("""case when count is not null
                                                                                     then (raw_rank + count - 1) else
                                                                                     raw_rank end""")))
    rankedNameGroup = rankedName
    n = df_w_rank.count()
    df_with_rank_groups = (df_w_rank.withColumn(rankedNameGroup,expr("""FLOOR({rankedName}
                                                                     *{k}/({n}+1))""".format(k=10, n=n,
                                                                     rankedName=rankedName))))
    df = df_with_rank_groups
aud_baskets_ranks = df_with_rank_groups.drop('raw_rank', 'count')
Related