Python Polars: применить функцию к двум столбцам и аргументу ⇐ Python

[Anonymous]

Введение
В Polars мне хотелось бы выполнять довольно сложные запросы и хотелось бы упростить процесс, разделив операции на методы. Прежде чем я смогу это сделать, мне нужно выяснить, как предоставить этой функции несколько столбцов и переменных.
Пример данных
# Libraries
import polars as pl
from datetime import datetime

# Data
test_data = pl.DataFrame({
"class": ['A', 'A', 'A', 'B', 'B', 'C'],
"date": [datetime(2020, 1, 31), datetime(2020, 2, 28), datetime(2021, 1, 31),
datetime(2022, 1, 31), datetime(2023, 2, 28),
datetime(2020, 1, 31)],
"status": [1,0,1,0,1,0]
})

Проблема
Для каждой группы я хотел бы знать, перекрывается ли справочная дата с годом-месяцем в столбце даты кадра данных. .
Я бы хотел сделать что-то подобное.
# Some date
reference_date = datetime(2020, 1, 2)

# What I would expect the query to look like
(test_data
.group_by("class")
.agg(
pl.col("status").count().alias("row_count"), #just to show code that works
pl.lit(reference_date).alias("reference_date"),
pl.col("date", "status")
.map_elements(lambda group: myfunc(group, reference_date))
.alias("point_in_time_status")
)
)

# The desired output
pl.DataFrame({
"class": ['A', 'B', 'C'],
"reference_date": [datetime(2020, 1, 2), datetime(2020, 1, 2), datetime(2020, 1, 2)],
"point_in_time_status": [1,0,0]
})

Но я просто не могу найти решения для выполнения операций над группами. Некоторые предлагают использовать pl.struct, но это просто выводит какой-то странный объект без столбцов или чего-либо, с чем можно работать.
Пример той же операции на R
# Loading library
library(tidyverse)

# Creating dataframe
df %
group_by("class") %>%
filter(format(date, "%Y-%m") == format(ref_date, "%Y-%m")) %>%
filter(status == 1)


Подробнее здесь: https://stackoverflow.com/questions/770 ... n-argument