Equivalence between shift() function (Python) and lead() function (R)

Question

General Question on understanding the SHIFT function in Python:

I am starting with this sample dataset:

start = structure(list(id = c(111L, 111L, 111L, 111L, 222L, 222L, 222L
), year = c(2010L, 2011L, 2012L, 2013L, 2018L, 2019L, 2020L),
col2 = c("A", "BB", "A", "C", "D", "EE", "F"), col3 = c(242L,
213L, 233L, 455L, 11L, 444L, 123L), col4 = c(1213L, 5959L,
9988L, 4242L, 333L, 1232L, 98L)), class = "data.frame", row.names = c(NA,
-7L))

This is what I want to do in R:

library(dplyr)

end <- start %>%
mutate(year_end = lead(year),
     col2_end = lead(col2),
     col3_end = lead(col3),
     col4_end = lead(col4)) %>%
mutate_at(vars(ends_with("_end")), ~ifelse(is.na(.), "END", .)) %>%
rename(year_start = year,
     col2_start = col2,
     col3_start = col3,
     col4_start = col4)

Now I am trying to do the same thing in Python.

I read that there is a SHIFT function in Python which is similar to the LEAD function in R - here is my attempt to recreate this work in Python:

import pandas as pd

start = pd.DataFrame({
    'id': [111, 111, 111, 111, 222, 222, 222],
    'year': [2010, 2011, 2012, 2013, 2018, 2019, 2020],
    'col2': ['A', 'BB', 'A', 'C', 'D', 'EE', 'F'],
    'col3': [242, 213, 233, 455, 11, 444, 123],
    'col4': [1213, 5959, 9988, 4242, 333, 1232, 98]
})

end = start.assign(
    year_end=start['year'].shift(-1),
    col2_end=start['col2'].shift(-1),
    col3_end=start['col3'].shift(-1),
    col4_end=start['col4'].shift(-1)
).fillna('END')

end = end.rename(columns={
    'year': 'year_start',
    'col2': 'col2_start',
    'col3': 'col3_start',
    'col4': 'col4_start'
})

I think this looks reasonable - but I was hoping to get a second set of eyes to verify my attempt. Any thoughts?

Are you in anyway renaming the columns? Because in R you are not. rename takes in (new = old) not (old = new) so Seems to me the rename has no effect. Check and you will note that the final data still has the initial names ie `year` etc unchanged — Onyambu, Apr 04 '23 at 16:30

score 3 · Accepted Answer · edited Apr 04 '23 at 17:46

In R:

start %>%
   mutate(across(-id, ~lead(as.character(.x),default = 'END'),.names = '{col}_end'))

   id year col2 col3 col4 year_end col2_end col3_end col4_end
1 111 2010    A  242 1213     2011       BB      213     5959
2 111 2011   BB  213 5959     2012        A      233     9988
3 111 2012    A  233 9988     2013        C      455     4242
4 111 2013    C  455 4242     2018        D       11      333
5 222 2018    D   11  333     2019       EE      444     1232
6 222 2019   EE  444 1232     2020        F      123       98
7 222 2020    F  123   98      END      END      END      END

In python:

(start.join(start.drop('id', axis = 1).shift(-1, fill_value = 'END')
    .add_suffix('_end')))
 
    id  year col2  col3  col4 year_end col2_end col3_end col4_end
0  111  2010    A   242  1213     2011       BB      213     5959
1  111  2011   BB   213  5959     2012        A      233     9988
2  111  2012    A   233  9988     2013        C      455     4242
3  111  2013    C   455  4242     2018        D       11      333
4  222  2018    D    11   333     2019       EE      444     1232
5  222  2019   EE   444  1232     2020        F      123       98
6  222  2020    F   123    98      END      END      END      END

Equivalence between shift() function (Python) and lead() function (R)

1 Answers1