基于历史行时间范围标记Flag column：3亿行DataFrame识别再入院

嘿，针对你这个处理3亿行患者住院数据的需求，我来分享几个高效的实现思路——毕竟数据量这么大，效率和内存控制绝对是关键！

首先明确核心需求：给已经按患者ID（Key）分组、且按入院日期（admit）升序排序的住院数据，标记符合两个条件的再入院记录：

• 条件1：当前行的Num1值在10-15之间（包含两端）
• 条件2：当前行的入院日期，落在该患者任意一条历史住院记录的「出院日期」到「出院后90天日期（90DayPostDischarge）」区间内

核心思路：利用有序数据优化计算 #

因为数据已经按患者ID分组、且入院日期升序，我们可以抓住这个特点，避免对每一行都遍历所有历史记录——这会让3亿行的计算慢到无法忍受。我们只需要维护一个「当前有效的出院区间列表」，只保留那些90天窗口还没过期的历史记录，每次只检查这些有效区间即可。

方案一：Pandas + Numba 加速（适合内存能容纳数据的场景） #

Numba可以把Python函数编译成机器码，处理大数组的速度比纯Pandas快很多，非常适合这种超大数据量的计算场景。

步骤1：准备数据（确保日期类型正确）

import pandas as pd
import numpy as np
from numba import jit

# 假设你的数据集叫df，先把日期列转为datetime类型，避免类型错误
df['admit'] = pd.to_datetime(df['admit'])
df['discharge_date'] = pd.to_datetime(df['discharge_date'])
df['90DayPostDischarge'] = pd.to_datetime(df['90DayPostDischarge'])

# 再次确认数据排序（如果之前的排序不严谨，补这一步）
df = df.sort_values(['Key', 'admit'])

步骤2：编写Numba加速的标记函数

@jit(nopython=True)
def flag_readmissions(admits, discharges, post_90, num1s):
    n = len(admits)
    flag = np.zeros(n, dtype=np.bool_)
    # 维护一个列表，存储还在90天窗口内的出院区间（出院日期, 出院后90天日期）
    active_intervals = []
    
    for i in range(n):
        current_admit = admits[i]
        current_num1 = num1s[i]
        
        # 先清理已经过期的区间（出院后90天早于当前入院日期的）
        while active_intervals and active_intervals[0][1] < current_admit:
            active_intervals.pop(0)
        
        # 检查两个条件
        condition1 = (10 <= current_num1 <= 15)
        condition2 = False
        # 遍历有效区间，看当前入院日期是否落在任意一个区间内
        for d, p in active_intervals:
            if d <= current_admit <= p:
                condition2 = True
                break
        
        flag[i] = condition1 and condition2
        
        # 将当前行的出院区间加入有效列表（因为入院日期升序，出院日期也必然升序，直接append即可）
        active_intervals.append((discharges[i], post_90[i]))
    
    return flag

这里把日期转成int64类型（timestamp的整数形式）是因为Numba对datetime类型的处理不如整数高效，后续应用时会自动转换。

步骤3：分组应用函数

# 对每个患者分组应用函数，注意把日期转成int64供Numba处理
df['Flag'] = df.groupby('Key').apply(
    lambda x: flag_readmissions(
        x['admit'].values.astype(np.int64),
        x['discharge_date'].values.astype(np.int64),
        x['90DayPostDischarge'].values.astype(np.int64),
        x['Num1'].values
    )
).explode().astype(int)  # 把布尔值转成0/1的整数标记，方便后续统计

方案二：Dask + Numba 处理超大数据（内存不够的场景） #

如果3亿行数据太大，无法一次性加载到内存，就用Dask来并行分块处理——它可以把数据拆成多个小分区，利用多核心计算，并且不需要把全量数据放在内存里。

import dask.dataframe as dd

# 把Pandas DataFrame转成Dask DataFrame，分区数根据你的内存大小调整（比如10-20个分区）
ddf = dd.from_pandas(df, npartitions=10)

# 定义Dask用的分组函数（直接调用之前写好的Numba函数）
def flag_readmissions_dask(x):
    admits = x['admit'].values.astype(np.int64)
    discharges = x['discharge_date'].values.astype(np.int64)
    post_90 = x['90DayPostDischarge'].values.astype(np.int64)
    num1s = x['Num1'].values
    return flag_readmissions(admits, discharges, post_90, num1s)

# 分组应用函数并计算最终结果
ddf['Flag'] = ddf.groupby('Key').apply(flag_readmissions_dask, meta=('Flag', 'bool')).astype(int)
result_df = ddf.compute()

注意事项 #

1. 先测小批量数据：先取100-1000个患者的子集测试函数是否正确，确认标记逻辑没问题再跑全量数据，避免浪费时间。
2. 日期边界调整：代码里默认包含区间端点（比如入院日期等于出院日期或出院后90天），如果业务上需要排除端点，修改条件里的</<=即可。
3. 内存监控：处理3亿行数据时，尽量关闭其他占用内存的程序，避免内存溢出。

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