Функция для добавления четвертей, представленных в виде строк в полярах
Основная идея состоит в том, чтобы взять строковые значения, представляющие конкатенацию год+квартал в формате 2020Q4, 1999Q1...ect, и добавить/вычесть произвольное количество кварталов к/из них. Эти строковые значения живут в фрейме данных «поляры» в столбце под названием «четверть». Я хочу создать еще один столбец «quarter_diff», в котором я вычитаю «четверть» предыдущей строки из значения текущей строки. Вычитание/сложение должно быть достаточно умным, чтобы увидеть, что если это 2020Q4, прибавив к нему 1 четверть, будет 2021Q1, а не 2020Q5
В идеале это должно быть сделано с помощью собственной функции «поляры», так как они уже оптимизированы для скорости.
Вот функция, написанная на питоне, поэтому не оптимизированная для скорости или распараллеливания поляров, и она работает только для добавления четвертей, а не для их вычитания.
def add_quarter(quarter_str):
year, quarter = quarter_str[:-2], quarter_str[-1]
if quarter == '4':
return f"{int(year) + 1}Q1"
else:
return f"{year}Q{int(quarter) + 1}"
Может ли кто-нибудь помочь в его рефакторинге как для сложения, так и для вычитания, а также для использования внутренней функции «поляры», например, offset_by() или другой?
Спасибо
Ответы 1
Использовать это
Это метод, который вы привязываете к пространству имен Expr, а затем можете использовать внутренние выражения. Он позволяет добавлять любое количество четвертей (я не тестировал отрицательные четверти, но должен работать).
def add_quarters(self,Quarters=1):
splt=self.str.split('Q')
current=splt.arr.first().cast(pl.Float64()) + ((splt.arr.last().cast(pl.Float64()))-1)/4
new_qt=current+Quarters/4
new_yr=new_qt.floor().cast(pl.Int32()).cast(pl.Utf8())
new_qtr=((new_qt*4).mod(4)+1).cast(pl.Int32()).cast(pl.Utf8())
return new_yr + "Q" + new_qtr
pl.Expr.add_quarters=add_quarters
тогда вы можете сделать
df.with_columns(next_q=pl.col('quarter').add_quarters(2))
shape: (3, 2)
┌─────────┬────────┐
│ quarter ┆ next_q │
│ --- ┆ --- │
│ str ┆ str │
╞═════════╪════════╡
│ 2020Q4 ┆ 2021Q2 │
│ 2021Q1 ┆ 2021Q3 │
│ 2020Q2 ┆ 2020Q4 │
└─────────┴────────┘
Это была моя первая мысль, оставляющая за собой мыслительный процесс/приложение
Это использует разделение вместо индекса строки символов, но тот же принцип, что и ваша функция python...
Если у вас есть
df=pl.DataFrame({'quarter':["2020Q4", "2021Q1", "2020Q2"]})
тогда вы можете сделать
df \
.with_columns(splt=pl.col("quarter").str.split('Q')) \
.with_columns(Year=pl.col('splt').arr.first().cast(pl.Int32()), Quarter=pl.col('splt').arr.last().cast(pl.Int32())) \
.drop('splt') \
.with_columns(next_quarter=pl.when(pl.col("Quarter")==4)
.then((pl.col('Year')+1).cast(pl.Utf8()) + "Q1")
.otherwise((pl.col('Year').cast(pl.Utf8())) + "Q" + (pl.col('Quarter')+1).cast(pl.Utf8()))) \
.drop(["Year",'Quarter'])
shape: (3, 2)
┌─────────┬──────────────┐
│ quarter ┆ next_quarter │
│ --- ┆ --- │
│ str ┆ str │
╞═════════╪══════════════╡
│ 2020Q4 ┆ 2021Q1 │
│ 2021Q1 ┆ 2021Q2 │
│ 2020Q2 ┆ 2020Q3 │
└─────────┴──────────────┘
Вторая мысль Вы можете объединить это в свое собственное выражение следующим образом...
def add_quarter(self):
splt=self.str.split('Q')
Year=splt.arr.first().cast(pl.Int32())
Quarter=splt.arr.last().cast(pl.Int32())
return pl.when(Quarter==4).then((Year+1).cast(pl.Utf8()) + "Q1").otherwise(Year.cast(pl.Utf8()) + "Q" + (Quarter+1).cast(pl.Utf8()))
pl.Expr.add_quarter=add_quarter
тогда просто...
df.with_columns(next_q=pl.col('quarter').add_quarter())
Спасибо. Таким образом, он правильно использует внутреннюю часть поляров. Хороший.