Python Polars Expression For Cumulative Row Count by Day by Category

Given a polars dataframe

import datetime
import polars as pl


df = pl.DataFrame(
    {
        "time": [
            datetime.datetime(2023, 1, 1, 9),
            datetime.datetime(2023, 1, 1, 10),
            datetime.datetime(2023, 1, 1, 12),
            datetime.datetime(2023, 1, 2, 9),
            datetime.datetime(2023, 1, 2, 10),
            datetime.datetime(2023, 1, 3, 12),
        ],
        "category": [1,1,2,1,2,1],
    }
)

I am seeking an expression row_count_by_day such that

expr = ...alias("row_count_by_day")
df = df.with_column(expr)
print(df)

yields

shape: (6, 3)
┌─────────────────────┬──────────┬──────────────────┐
│ time                ┆ category ┆ row_count_by_day │
│ ---                 ┆ ---      ┆ ---              │
│ datetime[μs]        ┆ i64      ┆ i64              │
╞═════════════════════╪══════════╪══════════════════╡
│ 2023-01-01 09:00:00 ┆ 1        ┆ 1                │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2023-01-01 10:00:00 ┆ 1        ┆ 2                │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2023-01-01 12:00:00 ┆ 2        ┆ 1                │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2023-01-02 09:00:00 ┆ 1        ┆ 1                │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2023-01-02 10:00:00 ┆ 2        ┆ 1                │
├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2023-01-03 12:00:00 ┆ 1        ┆ 1                │
└─────────────────────┴──────────┴──────────────────┘

That is, the cumulative row count for a particular category in a single day.

Ideally also I'd like to extend this to multiple categories. e.g. If category 1 has 5 rows with the same value in another column, they'd only be counted once.

I tried several combinations of window functions over category, but being a novice in polars, I didn't get close. My best guess was some sort of rolling_sum of a constant column with ones over the category with a window size of 1d. I couldn't get this to work.

It sounds like .dt.truncate is the piece you're missing.

df.with_columns(
   pl.first().cumcount()
     .over(pl.col("time").dt.truncate("1d"), "category")
     .alias("count") + 1
) 

shape: (6, 3)
┌─────────────────────┬──────────┬───────┐
│ time                ┆ category ┆ count │
│ ---                 ┆ ---      ┆ ---   │
│ datetime[μs]        ┆ i64      ┆ u32   │
╞═════════════════════╪══════════╪═══════╡
│ 2023-01-01 09:00:00 ┆ 1        ┆ 1     │
│ 2023-01-01 10:00:00 ┆ 1        ┆ 2     │
│ 2023-01-01 12:00:00 ┆ 2        ┆ 1     │
│ 2023-01-02 09:00:00 ┆ 1        ┆ 1     │
│ 2023-01-02 10:00:00 ┆ 2        ┆ 1     │
│ 2023-01-03 12:00:00 ┆ 1        ┆ 1     │
└─────────────────────┴──────────┴───────┘

You can do like this:

    df.with_columns(
    row_count_by_day=pl.repeat(1,pl.count())
    .cumsum()
    .over(
        pl.col('time')
        .dt.truncate("1d"),
        'category'))

The weird thing here is the pl.repeat(1,pl.count()) which seems like it's the same thing as pl.lit(1) but it's not because of the way polars broadcasts the scalar 1. If you do two contexts like this:

df.with_columns(one=pl.lit(1)).with_columns(
    row_count_by_day=pl.col('one')
    .cumsum()
    .over(
        pl.col('time')
        .dt.truncate("1d"),
        'category')).drop('one')

then the first with_columns broadcasts the 1 to all rows so then the second context works when referring to the one column.