将周对比报表代码从Pandas迁移至Polars遇阻，求解决方案

任职于教育科技公司，每周需生成当周与上周的对比报表，目前正将原有Pandas代码迁移至Polars，目标输出如下格式的字典：

{
"Admission/Retention": {
"Category": {
"1": {
"X": [
"2023-07-01",
"2023-07-02",
"2023-07-03",
"2023-07-04",
"2023-07-05",
"2023-07-06",
"2023-07-07"
],
"Y": [
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.5, 0.0, 0.0, 0.0, 0.0, 0.5, 0.0],
[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
]
},
"2": {
"X": [
"2023-07-01",
"2023-07-02",
"2023-07-03",
"2023-07-04",
"2023-07-05",
"2023-07-06",
"2023-07-07"
],
"Y": [
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
]
}
},
"Labels": ["Converted", "Converted_lastweek", "Converted_beginn_lastweek"]
}

但在Polars代码实现中遇到障碍，当前卡住的代码如下：

import polars as pl
import numpy as np
from datetime import date, datetime, timedelta

date = [
    "2023-06-24", "2023-06-25", "2023-06-26", "2023-06-27", "2023-06-28", "2023-06-29",
    "2023-06-30", "2023-07-01", "2023-07-02", "2023-07-03", "2023-07-04", "2023-07-05",
    "2023-07-06", "2023-07-07", "2023-06-24", "2023-06-25", "2023-06-26", "2023-06-27",
    "2023-06-28", "2023-06-29", "2023-06-30", "2023-07-01", "2023-07-02", "2023-07-03",
    "2023-07-04", "2023-07-05", "2023-07-06", "2023-07-07", "2023-06-24", "2023-06-25",
    "2023-06-26", "2023-06-27", "2023-06-28", "2023-06-29", "2023-06-30", "2023-07-01",
    "2023-07-02", "2023-07-03", "2023-07-04", "2023-07-05", "2023-07-06", "2023-07-07"
]

categories = [1] * 14 + [1] * 14 + [2] * 14
students = ["A"] * 14 + ["B"] * 14 + ["B"] * 14
fakes = [0] * 42
actives = [1] * 42

conversion_dates = ["6/23/2023", None, None, None, None, None, None, None, None, None, None, None, None, None,
                    None, None, None, None, "6/26/2023", None, None, None, None, None, None, None, None, None, None,
                    None, None, None, None, None, None, None, None, None, None, "7/3/2023", None, None]

converted = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]

df = pl.DataFrame({
    'date': date,
    'category': categories,
    'student': students,
    'fake': fakes,
    'active_student': actives,
    'conversion_date': conversion_dates,
    'converted': converted
})

def kpi_tracker(df):

    start_date = date(2023, 7, 1)
    dates = [start_date + timedelta(days=x) for x in range(0, 7)]
    past_week_start = start_date - timedelta(days=7)
    dates_last_week = [start_date + timedelta(days=x) for x in range(-7, 7)]

    bookings = np.zeros((1, len(dates)))
    bookings_lastweek = np.zeros((1, len(dates)))
    bookings_lastweek_beginning = np.zeros((1, len(dates)))

    mask = ((pl.col('converted') == 1) & (pl.col('conversion_date') < past_week_start))

    df = df.with_columns(
           pl.when(mask)
             .then(1)
             .alias('converted_stlw')
       )

#    temp = df[['date', 'converted', 'converted_stlw']][(df['active_student'])&(df['fake']==0)].groupby(['date']).mean()

#    temp = temp.reindex(dates_last_week)
#    temp['converted_stlw_shifted'] = temp['converted_stlw'].shift((7))
#    temp['converted_last_week'] = temp['converted'].shift((7))

#    temp = temp.reindex(dates)
#    converted = temp['converted'].fillna(0).tolist()
#    converted_lastweek = temp['converted_last_week'].fillna(0).tolist()
#    converted_beginn_lastweek = temp['converted_stlw_shifted'].fillna(0).tolist()

    result = {'X':dates,'Converted':converted, 'Converted_lastweek': converted_lastweek, 'Converted_beginn_lastweek': converted_beginn_lastweek}    
    return df

df_dataframe = pl.read_csv('data.csv')
df_dataframe = df_dataframe.with_columns(pl.col('conversion_date').str.strptime(pl.Date(), "%m/%d/%Y"))
df = df_dataframe.groupby('Category').apply(kpi_tracker)

核心问题分析 #

原代码的问题集中在：

• Polars无reindex方法，需用日期序列左连接补全缺失日期
• 日期列未转成Polars的Date类型，导致比较、分组出错
• kpi_tracker函数返回逻辑错误，未生成目标数据结构
• 分组应用后的结果未整理成目标字典格式

修正后的完整代码 #

import polars as pl
import numpy as np
from datetime import date, timedelta
import json

# 初始化测试数据（读取CSV可替换为pl.read_csv）
date_list = [
    "2023-06-24", "2023-06-25", "2023-06-26", "2023-06-27", "2023-06-28", "2023-06-29",
    "2023-06-30", "2023-07-01", "2023-07-02", "2023-07-03", "2023-07-04", "2023-07-05",
    "2023-07-06", "2023-07-07", "2023-06-24", "2023-06-25", "2023-06-26", "2023-06-27",
    "2023-06-28", "2023-06-29", "2023-06-30", "2023-07-01", "2023-07-02", "2023-07-03",
    "2023-07-04", "2023-07-05", "2023-07-06", "2023-07-07", "2023-06-24", "2023-06-25",
    "2023-06-26", "2023-06-27", "2023-06-28", "2023-06-29", "2023-06-30", "2023-07-01",
    "2023-07-02", "2023-07-03", "2023-07-04", "2023-07-05", "2023-07-06", "2023-07-07"
]

categories = [1] * 14 + [1] * 14 + [2] * 14
students = ["A"] * 14 + ["B"] * 14 + ["B"] * 14
fakes = [0] * 42
actives = [1] * 42

conversion_dates = ["6/23/2023", None, None, None, None, None, None, None, None, None, None, None, None, None,
                    None, None, None, None, "6/26/2023", None, None, None, None, None, None, None, None, None, None,
                    None, None, None, None, None, None, None, None, None, None, "7/3/2023", None, None]

converted = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0]

# 初始化DataFrame并统一日期类型
df = pl.DataFrame({
    'date': date_list,
    'category': categories,
    'student': students,
    'fake': fakes,
    'active_student': actives,
    'conversion_date': conversion_dates,
    'converted': converted
}).with_columns(
    pl.col('date').str.strptime(pl.Date(), "%Y-%m-%d"),
    pl.col('conversion_date').str.strptime(pl.Date(), "%m/%d/%Y").cast(pl.Date())
)

def kpi_tracker(df):
    start_date = date(2023, 7, 1)
    # 生成当周日期序列
    current_week_dates = pl.date_range(start_date, start_date + timedelta(days=6), interval="1d", eager=True)
    # 生成上周+本周的扩展日期序列（用于shift计算）
    extended_dates = pl.date_range(start_date - timedelta(days=7), start_date + timedelta(days=6), interval="1d", eager=True)
    
    # 过滤有效数据，标记上周初已转化用户，按日期求均值
    temp = df.filter(
        (pl.col('active_student') == 1) & (pl.col('fake') == 0)
    ).select('date', 'converted', 'conversion_date').with_columns(
        pl.when((pl.col('converted') == 1) & (pl.col('conversion_date') < start_date - timedelta(days=7)))
         .then(1)
         .otherwise(0)
         .alias('converted_stlw')
    ).group_by('date').mean()
    
    # 补全所有日期（替代Pandas的reindex）
    temp_full = pl.DataFrame({'date': extended_dates}).join(temp, on='date', how='left').fill_null(0)
    
    # 计算上周同期数据、上周初累计数据
    temp_full = temp_full.with_columns(
        pl.col('converted').shift(7).alias('converted_last_week'),
        pl.col('converted_stlw').shift(7).alias('converted_stlw_shifted')
    )
    
    # 筛选当周数据并整理格式
    current_week_data = temp_full.filter(pl.col('date').is_in(current_week_dates))
    return {
        "X": [d.strftime("%Y-%m-%d") for d in current_week_data['date'].to_list()],
        "Y": [
            current_week_data['converted'].to_list(),
            current_week_data['converted_last_week'].to_list(),
            current_week_data['converted_stlw_shifted'].to_list()
        ]
    }

# 按分类分组处理，生成分类结果字典
category_results = {}
for cat, group in df.group_by('category'):
    category_results[str(cat)] = kpi_tracker(group)

# 组装最终目标字典
final_result = {
    "Admission/Retention": {
        "Category": category_results,
        "Labels": ["Converted", "Converted_lastweek", "Converted_beginn_lastweek"]
    }
}

# 打印验证结果
print(json.dumps(final_result, indent=2))

关键修正点说明 #

1. 日期类型统一：将date和conversion_date转为Polars的Date类型，确保日期比较、分组逻辑正确。
2. 缺失日期补全：用pl.date_range生成完整日期序列，通过左连接join补全缺失日期，缺失值用fill_null(0)填充。
3. 分组逻辑优化：直接遍历分组结果，每个分类单独生成目标数据结构，避免groupby.apply的复杂返回处理。
4. 数据格式整理：将计算后的列转换为目标字典的X（日期字符串列表）和Y（三组数据列表）格式。

内容的提问来源于stack exchange，提问作者Rade