SQL to pandas: DENSE_RANK() OVER (PARTITION BY )

user_9066242 0
Viewed 600

I am trying to translate the following piece of SQL code to a pandas equivalent

SELECT
    t.company,
    t.topic,
    t.statement
FROM
    (
        SELECT
            e.company,
            e.topic,
            e.probability,
            e.distance,
            LOWER(e.statement) AS statement,
            dense_rank() OVER (PARTITION BY e.company,e.topic ORDER BY e.distance DESC) as rank
        FROM
            esg.group_dist e
    ) t
WHERE
    t.rank = 1
    AND t.topic IN ('green energy')
ORDER BY
    company,
    topic,
    rank

I got as far as

esg_group_dist['rank'] = esg_group_dist[['company', 'topic', 'probability', 'distance', 'sentence']] \
    .sort_values(by=['distance']) \
    .groupby(['company', 'topic']) \

I found the following SO thread that should contain a solution but I can't manage to successfully implement it for my usecase

Pandas DENSE RANK

Thanks!

pandas sql python
1 Answers

There is groupby.rank:

esg_group_dist['rank'] = (esg_group_dist.groupby(['company', 'topic'])
                             ['disance'].rank(method='dense', ascending=False)
                         )

However, looking at your entire query, it looks like you're trying to extract info where distance is maximum minimum within each group. You can do so faster with

(esg_group_dist[['company', 'topic', 'probability', 'distance', 'sentence']]
     .sort_values('distance')                            # sort values
     .drop_duplicates(['company','topic'], keep='last')  # keep the first rows
     .query('topic=="green energy"')                     # filter topic
)

Note: to find minimum rows, remove ascending=False and keep='last'. Also there is groupby().idxmin/idxmax() option`.

user_42384080
Related