pola-rs / polars

Dataframes powered by a multithreaded, vectorized query engine, written in Rust
https://docs.pola.rs
Other
29.19k stars 1.84k forks source link

Add .dt.replace #8879

Open mcrumiller opened 1 year ago

mcrumiller commented 1 year ago

Problem description

datetime.date has a replace function that lets you immediately set the day/month/year:

>>> from datetime import date
>>> date(2023, 1, 1).replace(month=3)
datetime.date(2023, 3, 1)

This would be a nice feature for series/expressions, especially if the input arg could also be an expression:

from datetime import date
import polars as pl
from polars import col

df = pl.DataFrame({
     "year": [2015, 2016, 2017],
    "my_date": [date(2023, 1, 1), date(2023, 1, 1), date(2023, 1, 1)],
})

df.select(
    col("my_date").dt.replace(year=col("year"))
)
shape: (3, 2)
┌──────┬────────────┐
│ year ┆ my_date    │
│ ---  ┆ ---        │
│ i64  ┆ date       │
╞══════╪════════════╡
│ 2015 ┆ 2015-01-01 │
│ 2016 ┆ 2016-01-01 │
│ 2017 ┆ 2017-01-01 │
└──────┴────────────┘
deanm0000 commented 1 year ago

I don't know how to monkey patch this into the dt namespace but otherwise, here ya go

def replace(self, year=None, month=None, day=None):
    """Mimic the datetime replace method"""
    if year is None and month is None and day is None:
        raise ValueError("don't you want to replace something, specify a year, month, and/or day")
    if year is None:
        outyear=self.dt.year()
    else:
        outyear=year
    if month is None:
        outmonth=self.dt.month()
    else:
        outmonth=month
    if day is None:
        outday=self.dt.day()
    else:
        outday=day
    return pl.date(outyear, outmonth, outday)
pl.Expr.replace=replace

then you can do

df.select(
    col("my_date").replace(year=col("year"))
)
mcrumiller commented 1 year ago

Thanks @deanm0000 . Yeah, I currently have my own monkeypatched methods but I felt it would be a good addition to the API.

FYI, when you post code in github, you can fence it with a language, e.g. ```python and you'll get syntax highlighting.

mcrumiller commented 2 weeks ago

@MarcoGorelli does someone want to accept this feature request? I can work on it but I don't want to spend the time if it's not desired/will be rejected.

mcrumiller commented 2 weeks ago

(FYI this is the original one :D)

MarcoGorelli commented 2 weeks ago

sure, marking as 'accepted', feel free to ping me for review