query-risk-user-features/scripts/run_social_features.py

#!/usr/bin/env python3
"""Query and aggregate share relationships and source-group features."""

from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd


PROJECT_DIR = Path(__file__).resolve().parents[1]
WORKSPACE_ROOT = Path(__file__).resolve().parents[2]
for path in (PROJECT_DIR, WORKSPACE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from odps_module import ODPSClient  # noqa: E402
from src.social_logic import (  # noqa: E402
    INVALID_GROUP_IDS,
    aggregate_groups,
    aggregate_relationships,
    build_group_sql,
    build_relationship_sql,
    latest_partition,
)
from src.user_input import load_users  # noqa: E402


RELATIONSHIP_COLUMNS = [
    "sample_mid", "user_type", "sub_category", "anchor_dt",
    "window_start_dt", "window_end_dt", "click_event_id", "click_dt",
    "click_ts", "shareid", "clickobjectid", "hotsencetype", "valid_opengid",
    "sessionid", "subsessionid", "share_dt", "share_ts", "source_share_mid",
    "is_self_source", "source_is_blacklist", "source_risk_type_list",
    "share_click_gap_seconds", "previous_clickobjectid", "previous_click_ts",
    "is_consecutive_same_card",
]

GROUP_COLUMNS = [
    "sample_mid", "user_type", "sub_category", "window_start_dt",
    "window_end_dt", "opengid", "group_path_pv", "group_visit_uv",
    "group_risk_user_uv", "group_risk_user_rate", "group_source_sharer_uv",
    "group_risk_source_sharer_uv", "group_risk_source_sharer_rate",
    "group_real_play_pv", "group_real_play_uv", "group_real_play_user_rate",
    "group_real_play_per_visitor",
]


def share_search_start(value: str) -> str:
    date = datetime.strptime(value, "%Y%m%d") - timedelta(days=180)
    return date.strftime("%Y%m%d")


def build_group_input(users: pd.DataFrame, detail: pd.DataFrame) -> pd.DataFrame:
    required = {"用户id", "来源群id"}
    missing = required - set(detail.columns)
    if missing:
        raise ValueError("基础明细缺少字段: " + ", ".join(sorted(missing)))
    windows = users.set_index("mid")
    rows = []
    for record in detail[["用户id", "来源群id"]].to_dict("records"):
        mid = str(record["用户id"])
        if mid not in windows.index:
            raise ValueError(f"基础明细用户不在输入人员表: {mid}")
        group_ids = {
            value.strip()
            for value in str(record["来源群id"] or "").split(",")
            if value.strip() not in INVALID_GROUP_IDS
        }
        for opengid in sorted(group_ids):
            rows.append(
                {
                    "sample_mid": mid,
                    "user_type": windows.at[mid, "user_type"],
                    "sub_category": windows.at[mid, "sub_category"],
                    "window_start_dt": windows.at[mid, "window_start_dt"],
                    "window_end_dt": windows.at[mid, "window_end_dt"],
                    "opengid": opengid,
                }
            )
    return pd.DataFrame(
        rows,
        columns=[
            "sample_mid", "user_type", "sub_category", "window_start_dt",
            "window_end_dt", "opengid",
        ],
    ).drop_duplicates(["sample_mid", "opengid"])


def read_instance(instance: object) -> pd.DataFrame:
    instance.wait_for_success()
    with instance.open_reader(tunnel=True) as reader:
        frame = reader.to_pandas()
    frame.columns = [str(column).lower() for column in frame.columns]
    return frame


def write_outputs(
    users: pd.DataFrame,
    behavior_detail: pd.DataFrame,
    relationships: pd.DataFrame,
    group_detail: pd.DataFrame,
    output_dir: Path,
) -> None:
    user_relationship = aggregate_relationships(users, relationships)
    user_group = aggregate_groups(users, group_detail)
    social = user_relationship.merge(user_group, on="用户id", how="left", validate="one_to_one")
    extended = behavior_detail.merge(social, on="用户id", how="left", validate="one_to_one")

    order = {mid: index for index, mid in enumerate(users["mid"])}
    for frame, column in ((social, "用户id"), (extended, "用户id")):
        frame["_input_order"] = frame[column].map(order)
        frame.sort_values("_input_order", kind="stable", inplace=True)
        frame.drop(columns="_input_order", inplace=True)

    output_dir.mkdir(parents=True, exist_ok=True)
    social.to_csv(output_dir / "user_social_features.csv", index=False)
    relationships.to_csv(output_dir / "relationship_detail.csv", index=False)
    group_detail.to_csv(output_dir / "group_detail.csv", index=False)
    extended.to_csv(output_dir / "user_detail_with_social.csv", index=False)
    with pd.ExcelWriter(output_dir / "social_features.xlsx", engine="openpyxl") as writer:
        extended.to_excel(writer, sheet_name="用户扩展明细", index=False)
        social.to_excel(writer, sheet_name="用户社交特征", index=False)
        group_detail.to_excel(writer, sheet_name="来源群明细", index=False)
        relationships.to_excel(writer, sheet_name="分享关系明细", index=False)

    expected = len(users)
    if len(social) != expected or social["用户id"].nunique() != expected:
        raise AssertionError("社交特征未保持输入用户粒度")
    if len(extended) != expected or extended["用户id"].nunique() != expected:
        raise AssertionError("扩展明细未保持输入用户粒度")
    if not group_detail.empty and group_detail["opengid"].isin(INVALID_GROUP_IDS).any():
        raise AssertionError("群级结果包含无效群ID")
    gaps = pd.to_numeric(relationships["share_click_gap_seconds"], errors="coerce").dropna()
    if (gaps < 0).any():
        raise AssertionError("存在负分享点击时间差")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--behavior-detail", type=Path, required=True)
    parser.add_argument("--run-name", required=True)
    parser.add_argument("--lookback-days", type=int, default=180)
    parser.add_argument("--blacklist-dt", default="")
    parser.add_argument("--sql-only", action="store_true")
    parser.add_argument("--finalize-only", action="store_true")
    parser.add_argument("--collect-existing", action="store_true")
    parser.add_argument("--modules", default="")
    args = parser.parse_args()

    users = load_users(args.input, lookback_days=args.lookback_days)
    users["share_search_start_dt"] = users["window_start_dt"].map(share_search_start)
    behavior_detail = pd.read_csv(args.behavior_detail, dtype=str).fillna("")
    if set(behavior_detail["用户id"]) != set(users["mid"]):
        raise AssertionError("基础明细与输入人员的MID集合不一致")
    group_input = build_group_input(users, behavior_detail)

    output_dir = PROJECT_DIR / "output" / args.run_name / "social"
    raw_dir = output_dir / "raw"
    sql_dir = PROJECT_DIR / "sql" / args.run_name / "social"
    for directory in (output_dir, raw_dir, sql_dir):
        directory.mkdir(parents=True, exist_ok=True)
    users.to_csv(output_dir / "input_users.csv", index=False)
    group_input.to_csv(output_dir / "input_user_groups.csv", index=False)

    if args.finalize_only:
        relationships = pd.read_csv(raw_dir / "relationships.csv")
        groups = (
            pd.read_csv(raw_dir / "groups.csv")
            if (raw_dir / "groups.csv").exists()
            else pd.DataFrame(columns=GROUP_COLUMNS)
        )
        write_outputs(users, behavior_detail, relationships, groups, output_dir)
        print(f"output={output_dir / 'social_features.xlsx'}")
        return

    client = ODPSClient(project="loghubods")
    blacklist_dt = args.blacklist_dt or latest_partition(client, "ad_report_user_blacklist")
    queries = {"relationships": build_relationship_sql(users, blacklist_dt)}
    if not group_input.empty:
        queries["groups"] = build_group_sql(group_input, blacklist_dt)
    for name, sql in queries.items():
        (sql_dir / f"{name}.sql").write_text(sql + "\n", encoding="utf-8")
    (output_dir / "query_metadata.json").write_text(
        json.dumps(
            {
                "blacklist_dt": blacklist_dt,
                "users": len(users),
                "user_group_edges": len(group_input),
                "distinct_groups": group_input["opengid"].nunique() if len(group_input) else 0,
            },
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    if args.sql_only:
        print(f"sql_dir={sql_dir}")
        return

    selected = [item.strip() for item in args.modules.split(",") if item.strip()]
    unknown = sorted(set(selected) - set(queries))
    if unknown:
        raise ValueError("未知模块: " + ", ".join(unknown))
    selected_queries = {name: queries[name] for name in selected} if selected else queries
    runs_path = output_dir / "odps_runs.json"

    if args.collect_existing:
        history = json.loads(runs_path.read_text(encoding="utf-8"))
        latest = {item["name"]: item for item in history}
        instances = {
            name: client.odps.get_instance(latest[name]["instance_id"])
            for name in selected_queries
        }
    else:
        instances = {}
        new_runs = []
        for name, sql in selected_queries.items():
            instance = client.odps.run_sql(sql, hints={"odps.sql.allow.cartesian": "true"})
            instances[name] = instance
            run = {
                "name": name,
                "instance_id": instance.id,
                "logview": instance.get_logview_address(),
            }
            new_runs.append(run)
            print(f"[{name}] instance={run['instance_id']}", flush=True)
            print(f"[{name}] logview={run['logview']}", flush=True)
        history = json.loads(runs_path.read_text(encoding="utf-8")) if runs_path.exists() else []
        runs_path.write_text(
            json.dumps([*history, *new_runs], ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )

    results = {}
    for name, instance in instances.items():
        frame = read_instance(instance)
        frame.to_csv(raw_dir / f"{name}.csv", index=False)
        results[name] = frame
        print(f"[{name}] rows={len(frame)}", flush=True)
    for name in queries:
        if name not in results:
            path = raw_dir / f"{name}.csv"
            if not path.exists():
                raise ValueError(f"缺少模块结果: {path}")
            results[name] = pd.read_csv(path)
    groups = results.get("groups", pd.DataFrame(columns=GROUP_COLUMNS))
    write_outputs(users, behavior_detail, results["relationships"], groups, output_dir)
    print(f"output={output_dir / 'social_features.xlsx'}")


if __name__ == "__main__":
    main()