jiangcang_vj/result_keep_verify.py

import os
import datetime
import pandas as pd
from data_loader import mongo_con_parse, validate_keep_one_line, fill_hourly_crawl_date
from config import vj_flight_route_list_hot, vj_flight_route_list_nothot, \
    CLEAN_VJ_HOT_NEAR_INFO_TAB, CLEAN_VJ_HOT_FAR_INFO_TAB, CLEAN_VJ_NOTHOT_NEAR_INFO_TAB, CLEAN_VJ_NOTHOT_FAR_INFO_TAB


def _validate_keep_info_df(df_keep_info_part):
    client, db = mongo_con_parse()
    count = 0

    if "price_diff" not in df_keep_info_part.columns:
        df_keep_info_part["price_diff"] = 0
    if "time_diff_hours" not in df_keep_info_part.columns:
        df_keep_info_part["time_diff_hours"] = 0

    for idx, row in df_keep_info_part.iterrows():
        df_keep_info_part.at[idx, "price_diff"] = 0
        df_keep_info_part.at[idx, "time_diff_hours"] = 0

        city_pair = row['city_pair']
        flight_day = row['flight_day']
        flight_number_1 = row['flight_number_1']
        flight_number_2 = row['flight_number_2']
        baggage = row['baggage']
        update_hour = row['update_hour']
        update_dt = pd.to_datetime(update_hour, format='%Y-%m-%d %H:%M:%S')
        into_update_hour = row['into_update_hour']
        into_update_dt = pd.to_datetime(into_update_hour, format='%Y-%m-%d %H:%M:%S')

        entry_price = pd.to_numeric(row.get('adult_total_price'), errors='coerce')
        if city_pair in vj_flight_route_list_hot:
            table_name_far = CLEAN_VJ_HOT_FAR_INFO_TAB
            table_name_near = CLEAN_VJ_HOT_NEAR_INFO_TAB
        elif city_pair in vj_flight_route_list_nothot: 
            table_name_far = CLEAN_VJ_NOTHOT_FAR_INFO_TAB
            table_name_near = CLEAN_VJ_NOTHOT_NEAR_INFO_TAB

        # 分别从远期表和近期表里查询
        df_query_far = validate_keep_one_line(db, table_name_far, city_pair, flight_day, flight_number_1, flight_number_2, baggage, into_update_hour)
        df_query_near = validate_keep_one_line(db, table_name_near, city_pair, flight_day, flight_number_1, flight_number_2, baggage, into_update_hour)
        # 合并
        df_query = pd.concat([df_query_far, df_query_near]).reset_index(drop=True)
        if (not df_query.empty) and pd.notna(entry_price):
            if ("adult_total_price" in df_query.columns) and ("crawl_date" in df_query.columns):
                df_query["adult_total_price"] = pd.to_numeric(df_query["adult_total_price"], errors="coerce")
                df_query["crawl_dt"] = pd.to_datetime(df_query["crawl_date"], errors="coerce")
                df_query = (
                    df_query.dropna(subset=["adult_total_price", "crawl_dt"])
                    .sort_values("crawl_dt")
                    .reset_index(drop=True)
                )
                mask_drop = df_query["adult_total_price"] < entry_price
                if mask_drop.any():
                    first_row = df_query.loc[mask_drop].iloc[0]
                    price_diff = entry_price - first_row["adult_total_price"]
                    time_diff_hours = (first_row["crawl_dt"] - into_update_dt) / pd.Timedelta(hours=1)
                    df_keep_info_part.at[idx, "price_diff"] = round(float(price_diff), 2)
                    df_keep_info_part.at[idx, "time_diff_hours"] = round(float(time_diff_hours), 2)
                    pass
        
        del df_query
        del df_query_far
        del df_query_near

        count += 1
        if count % 5 == 0:
            print(f"cal count: {count}")
    
    print(f"计算结束")
    client.close()

    return df_keep_info_part


def verify_process(min_batch_time_str, max_batch_time_str):
    object_dir = "./keep_0"

    output_dir = f"./validate/keep"
    os.makedirs(output_dir, exist_ok=True)

    timestamp_str = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    save_scv = f"result_keep_verify_{timestamp_str}.csv"
    output_path = os.path.join(output_dir, save_scv)
    
    # 获取今天的日期
    # today_str = pd.Timestamp.now().strftime('%Y-%m-%d')

    # 检查目录是否存在
    if not os.path.exists(object_dir):
        print(f"目录不存在: {object_dir}")
        return
                    
    # 获取所有以 keep_info_ 开头的 CSV 文件
    csv_files = []
    for file in os.listdir(object_dir):
        if file.startswith("keep_info_") and file.endswith(".csv"):
            csv_files.append(file)

    if not csv_files:
        print(f"在 {object_dir} 中没有找到 keep_info_ 开头的 CSV 文件")
        return
    
    csv_files.sort()
    # print(csv_files)

    min_batch_dt = datetime.datetime.strptime(min_batch_time_str, "%Y%m%d%H%M")
    min_batch_dt = min_batch_dt.replace(minute=0, second=0, microsecond=0)
    max_batch_dt = datetime.datetime.strptime(max_batch_time_str, "%Y%m%d%H%M")
    max_batch_dt = max_batch_dt.replace(minute=0, second=0, microsecond=0)

    if min_batch_dt is not None and max_batch_dt is not None and min_batch_dt > max_batch_dt:
        print(f"时间范围非法: min_batch_time_str({min_batch_time_str}) > max_batch_time_str({max_batch_time_str})，退出")
        return
    
    # 从所有的 keep_info 文件中
    for csv_file in csv_files:
        batch_time_str = (
            csv_file.replace("keep_info_", "").replace(".csv", "")
        )
        batch_dt = datetime.datetime.strptime(batch_time_str, "%Y%m%d%H%M")
        batch_hour_dt = batch_dt.replace(minute=0, second=0, microsecond=0)
        
        if min_batch_dt is not None and batch_hour_dt < min_batch_dt:
            continue
        if max_batch_dt is not None and batch_hour_dt > max_batch_dt:
            continue

        # 读取 CSV 文件
        csv_path = os.path.join(object_dir, csv_file)
        try:
            df_keep_info = pd.read_csv(csv_path)
        except Exception as e:
            print(f"read {csv_path} error: {str(e)}")
            df_keep_info = pd.DataFrame()
        
        if df_keep_info.empty:
            print(f"keep_info数据为空: {csv_file}")
            continue
        
        df_keep_info_del = df_keep_info[df_keep_info['keep_flag'] == -1].reset_index(drop=True)
        df_keep_info_del = _validate_keep_info_df(df_keep_info_del)
        df_keep_info_del['del_batch_time_str'] = batch_time_str

        write_header = not os.path.exists(output_path)
        df_keep_info_del.to_csv(output_path, mode="a", header=write_header, index=False, encoding="utf-8-sig")
        del df_keep_info_del
        print(f"批次:{batch_time_str} 检验结束")

    print("检验结束")
    print()
        

if __name__ == "__main__":
    verify_process("202603121800", "202603131600")
    pass