Writing advanced SQL queries in pandas | by Zolzaya Luvsandorj | Towards Data Science

Being able to skillfully manipulate data with both SQL and pandas, a data analysis library in Python, is a valuable skill to have for data analysts, data scientists and anyone working with data. In this post, we will look at selected few advanced SQL queries and their counterparts in pandas.
If you are still new to pandas, you may want to check out 10 minutes to pandas or my previous post where I covered side-by-side comparison of basic queries and their equivalents in pandas. This post assumes that the reader is comfortable with querying in SQL.
We will create a small dataset to use. Let’s assume we have two imaginary people’s trip data from the past 2 years:
df = pd.DataFrame({'name': ['Ann', 'Ann', 'Ann', 'Bob', 'Bob'],                    'destination': ['Japan', 'Korea', 'Switzerland',                                    'USA', 'Switzerland'],                    'dep_date': ['2019-02-02', '2019-01-01',                                 '2020-01-11', '2019-05-05',                                 '2020-01-11'],                    'duration': [7, 21, 14, 10, 14]})df
Let’s define dep_date as the departure date to the destination and assume that duration shows trip length in days.
We will start with the basic versions of LEAD() and LAG(). For each trip, let’s pull the departure date of the next trip: lead1, the second next trip: lead2, the previous trip: lag1 and the third previous trip: lag3.
SELECT name       , destination       , dep_date       , duration       , LEAD(dep_date) OVER(ORDER BY dep_date, name) AS lead1       , LEAD(dep_date, 2) OVER(ORDER BY dep_date, name) AS lead2       , LAG(dep_date) OVER(ORDER BY dep_date, name) AS lag1       , LAG(dep_date, 3) OVER(ORDER BY dep_date, name) AS lag3FROM df
To get the same output in pandas, we use shift():
df.sort_values(['dep_date', 'name'], inplace=True)df.assign(lead1 = df['dep_date'].shift(-1),          lead2 = df['dep_date'].shift(-2),          lag1 = df['dep_date'].shift(),          lag3 = df['dep_date'].shift(3))
In the first line, we sorted the data with sort_values() once in place since we are using the same logic for all shifting operations. Doing it this way is more efficient than sorting the data every time we create a new column like below:
df.assign(lead1 = df.sort_values(['dep_date', 'name'])['dep_date']                    .shift(-1),          lead2 = df.sort_values(['dep_date', 'name'])['dep_date']                    .shift(-2),          lag1 = df.sort_values(['dep_date', 'name'])['dep_date']                   .shift(),          lag3 = df.sort_values(['dep_date', 'name'])['dep_date']                   .shift(3))
Output from this inefficient approach will contain the right values but the order of the rows will be same as the original data and will not be sorted.
Now let’s look at the remaining 4 lines of the code. In pandas, we can create multiple columns with assign(). However, the new columns haven’t been added to the DataFrame. If we wanted to add the new columns to df , we will need to assign like this:
df.sort_values(['dep_date', 'name'], inplace=True)df = df.assign(lead1 = df['dep_date'].shift(-1),               lead2 = df['dep_date'].shift(-2),               lag1 = df['dep_date'].shift(),               lag3 = df['dep_date'].shift(3))
We will look at another example for LEAD() with PARTITION BY() in the last section. Now, let’s look at how to manipulate date/datetime columns.
In this section, we will extract day name of the week from departure date: day, name of the month from departure date: month, number of days elapsed since previous trip: diff and arrival date: arr_date.
SELECT name       , destination       , dep_date       , duration       , DATENAME(WEEKDAY, dep_date) AS day       , DATENAME(MONTH, dep_date) AS month       , DATEDIFF(DAY,                    LAG(dep_date) OVER(ORDER BY dep_date, name),                   dep_date) AS diff       , DATEADD(DAY, day, dep_date) AS arr_dateFROM df
Firstly, we must ensure the column data types are the right ones:
# Convert to proper dtypedf['dep_date'] = pd.to_datetime(df['dep_date'])df['duration'] = pd.to_timedelta(df['duration'], 'D')
Converting dep_date to datetime will allow us to access many date parts with the .dt accessor. For instance: df[‘dep_date’].dt.year will give us the year (equivalent to DATEPART(YEAR, dep_date) in SQL).
Converting duration to timedelta allows us to add it to a datetime column to get another datetime column.
Having done the data type conversion, let’s see the comparison:
df.sort_values(['dep_date', 'name'], inplace=True)df.assign(day = df['dep_date'].dt.day_name(),          month = df['dep_date'].dt.month_name(),          diff = df['dep_date'] - df['dep_date'].shift(),          arr_date = df['dep_date'] + df['duration'])
As you can see, the manipulation is simple once the data types are correctly set. Of note, we can check the data types of columns with the dtypes attribute of a DataFrame. Here’s an example:
df.sort_values(['dep_date', 'name'], inplace=True)df.assign(day = df['dep_date'].dt.day_name(),          month = df['dep_date'].dt.month_name(),          diff = df['dep_date'] - df['dep_date'].shift(),          arr_date = df['dep_date'] + df['duration']).dtypes
In pandas, when finding the difference between two dates, it returns a timedelta column. Therefore, we can see that column diff is actually a timedelta.
You may have used at least one of these functions before in SQL. The key difference between these three is their treatment in ranking ties and ranks thereafter. Let’s create three columns that rank each record based on duration using each of these method: row_number_d, rank_d and dense_rank_d.
When using ROW_NUMBER(), ties are given a different rank unlike the other two. To control how ties are treated with ROW_NUMBER(), we used name to breaks ties.
SELECT name       , destination       , dep_date       , duration       , ROW_NUMBER() OVER(ORDER BY duration, name) AS row_number_d       , RANK() OVER(ORDER BY duration) AS rank_d       , DENSE_RANK() OVER(ORDER BY duration) AS dense_rank_dFROM df
The equivalent of the query in pandas is possible with rank():
df.sort_values(['duration', 'name']).assign(    row_number_d = df['duration'].rank(method='first').astype(int),    rank_d = df['duration'].rank(method='min').astype(int),    dense_rank_d = df['duration'].rank(method='dense').astype(int))
We specify the relevant method inside rank() as illustrated above. Of note, there are other variations we haven’t covered here. If you are interested to find out, have a look at the documentation.
We did something slightly different with this example. Unlike examples in previous two sections, we didn’t sort the data in place. Instead, we chained methods and haven’t modified the data at all.
You may have also noticed that we have included astype(int). This is to convert the ranks in floats to integers. If you prefer to see floats, then that portion of the code could be removed:
df.sort_values(['duration', 'name']).assign(    row_number_d = df['duration'].rank(method='first'),    rank_d = df['duration'].rank(method='min'),    dense_rank_d = df['duration'].rank(method='dense'))
In this section, we will create column for the longest trip duration: max_dur, total duration for all trips: sum_dur, average duration of trip per person: avg_dur_name and cumulative sum of trip length for each person: cum_sum_dur_name.
SELECT name       , destination       , dep_date        , duration       , MAX(duration) OVER() AS max_dur       , SUM(duration) OVER() AS sum_dur       , AVG(duration) OVER(PARTITION BY name) AS avg_dur_name       , SUM(duration) OVER(PARTITION BY name ORDER BY dep_date                            RANGE BETWEEN UNBOUNDED PRECEDING                            AND CURRENT ROW) AS cum_sum_dur_nameFROM df
This can be done like this in pandas:
df.assign(max_dur=df['duration'].max(),          sum_dur=df['duration'].sum(),          avg_dur_name=df.groupby('name')['duration']                         .transform('mean'),          cum_sum_dur_name=df.sort_values('dep_date')                             .groupby('name')['duration']                             .transform('cumsum'))
Appending aggregate statistics such as max_dur and sum_dur is simple in pandas because of broadcasting. Essentially, if we try to assign a scalar value to a new column in pandas, the value is broadcasted across all rows. PARTITION BY is achieved with combination of groupby() and transform().
In this output, the data is in the same order as before (i.e. unsorted) whereas it will be sorted by dep_date for each person in SQL. We have only sorted the data when creating cum_sum_dur_name. If we desire the output to be sorted, the code changes to:
df.sort_values(['name', 'dep_date'], inplace=True)df.assign(max_dur=df['duration'].max(),          sum_dur=df['duration'].sum(),          avg_dur_name=df.groupby('name')['duration']                         .transform('mean'),          cum_sum_dur_name=df.groupby('name')['duration']                             .transform('cumsum'))
This will be identical to the SQL output.
Lastly, let’s write one query that covers pieces of the four different areas we just visited. Here’s our last example:
SELECT ROW_NUMBER() OVER(PARTITION BY name                          ORDER BY dep_date) AS number       , name       , destination       , DATENAME(MONTH, dep_date) AS dep_month       , dep_date       , DATEADD(DAY, duration, dep_date) AS arr_date       , LEAD(dep_date) OVER(PARTITION BY NAME                              ORDER BY dep_date) AS next_dep_date       , DATEDIFF(DAY,                   dep_date,                   LEAD(dep_date) OVER(PARTITION BY NAME                                       ORDER BY dep_date)) AS gap       , duration       , AVG(1.0 * duration) OVER() AS avg_dur       , AVG(1.0 * duration) OVER(PARTITION BY name) AS avg_dur_name       , SUM(duration) OVER(PARTITION BY name ORDER BY dep_date                            RANGE BETWEEN UNBOUNDED PRECEDING                             AND CURRENT ROW) AS cum_sum_dur_name   FROM dfORDER BY name, dep_date
I believe we learn more when we practice ourselves. I encourage you to try translating this SQL query to pandas on your own before seeing the comparison that we are about to look at. 👀
Here is the comparison:
# Convert to proper typedf['dep_date'] = pd.to_datetime(df['dep_date'])df['duration'] = pd.to_timedelta(df['duration'], 'D')# Sort datadf.sort_values(['name', 'dep_date'], inplace=True)# Append new columns to datadf = df.assign(number=df.groupby('name')['dep_date']                        .rank('min')                        .astype(int),               dep_month = df['dep_date'].dt.month_name(),               arr_date = df['dep_date'] + df['duration'],               next_dep_date = df.groupby('name')['dep_date']                                 .transform(lambda x: x.shift(-1)),               gap = df.groupby('name')['dep_date']                       .transform(lambda x: x.shift(-1))-                                            df['dep_date'],               avg_dur = df['duration'].mean(),               avg_dur_name = df.groupby('name')['duration']                                .transform(lambda x: x.mean()),               cum_sum_dur_name = df.groupby('name')['duration']                                   .transform(lambda x: x.cumsum()))# Reorder columnscolumns = ['number', 'name', 'destination', 'dep_month',            'dep_date', 'arr_date', 'next_dep_date',            'gap', 'duration', 'avg_dur', 'avg_dur_name',            'cum_sum_dur_name']df[columns]
We have converted data types, sorted the data in place and modified it with new columns. To view the data with columns ordered in the same way as SQL query, we used a list which outlines the order of the columns. If you only run df, the columns would be ordered differently.
You may have noticed that avg_dur is expressed in combination of days, hours and minutes: ‘13 days 04:48:00’. This is equivalent to 13 + (4*60+48) / (24*60) = 13.2 days. If we wanted to see just 13.2, then we could replace df[‘duration’].mean() with this snippet which converts duration to numeric type before finding mean: pd.to_numeric(df[‘duration’].dt.days).mean().
When creating avg_dur_name and cum_sum_dur_name, we used lambda function instead of using a syntax like this:
avg_dur_name = df.groupby('name')['duration'].transform('mean'),cum_sum_dur_name = df.groupby('name')['duration']                     .transform('cumsum')
That’s because if we try to do that we will come across an issue described here. Therefore, we used one of the work-arounds instead.
Voila❕ This remarks the end of our comparisons.
Would you like to access more content like this? Medium members get unlimited access to any articles on Medium. If you become a member using my referral link, a portion of your membership fee will directly go to support me.
Thank you for reading my post. Hope you find this post useful ✂️ and learned a bit more about pandas. If you are interested to learn more about pandas, here are links to some of my posts:
◼️️ 5 tips for data aggregation in pandas◼️️ Writing 5 common SQL queries in pandas◼️️ 5 tips for pandas users◼️️ How to transform variables in a pandas DataFrame
Bye for now 🏃💨