Agent/examples/auto_put_ad_mini/tools/video_feature_query.py

"""落地页视频素材召回使用的 ODPS 视频特征查询。

videoContentList 只负责选择落地页视频；素材召回特征按视频 ID 从
loghubods.dwd_video_element_contribution_analysis 读取。
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

logger = logging.getLogger(__name__)

_MINI_DIR = Path(__file__).resolve().parent.parent


@dataclass(frozen=True)
class VideoElementFeature:
    video_id: int
    element_dimension: str
    point_type: str
    standard_element: str
    contribution_score: float
    dt: str


CREATE_CACHE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS video_element_feature_cache (
    id BIGINT AUTO_INCREMENT PRIMARY KEY COMMENT '自增主键',
    video_id BIGINT NOT NULL COMMENT '业务视频ID，对应 ODPS vid',
    dt VARCHAR(8) NOT NULL COMMENT 'ODPS 分区日期 YYYYMMDD',
    point_type VARCHAR(50) NOT NULL COMMENT '点类型',
    standard_element VARCHAR(1024) NOT NULL COMMENT '标准化元素或解构选题文本',
    element_dimension VARCHAR(50) NOT NULL DEFAULT '实质' COMMENT '元素维度',
    contribution_score DOUBLE NOT NULL DEFAULT 0 COMMENT '贡献分',
    is_miss BOOLEAN NOT NULL DEFAULT FALSE COMMENT '是否为无特征负缓存',
    raw_element VARCHAR(255) DEFAULT NULL COMMENT '原始元素名称',
    element_id VARCHAR(100) DEFAULT NULL COMMENT '标准化元素ID',
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
    UNIQUE KEY uk_dt_video_point_element (dt, video_id, element_dimension, point_type, standard_element(191)),
    KEY idx_video_dt (video_id, dt),
    KEY idx_dt (dt),
    KEY idx_updated_at (updated_at)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='视频元素特征缓存'
"""


def _quote_sql_string(value: str) -> str:
    """把字符串转义并加单引号,用于拼 ODPS SQL 字面量。"""
    return "'" + value.replace("\\", "\\\\").replace("'", "\\'") + "'"


def _chunks(items: list[int], size: int) -> Iterable[list[int]]:
    """把列表按 size 切成批次,逐批 yield。"""
    for i in range(0, len(items), size):
        yield items[i:i + size]


def _get_odps_client():
    """获取 loghubods 项目的 ODPS client(动态注入 tools 目录到 sys.path)。"""
    import sys

    tools_dir = _MINI_DIR / "tools"
    if str(tools_dir) not in sys.path:
        sys.path.insert(0, str(tools_dir))
    from odps_module import get_odps_client

    return get_odps_client(project="loghubods")


def _latest_partition_dt(client) -> str:
    """获取 dwd_video_element_contribution_analysis 表的最新分区 dt(YYYYMMDD)。

    先走表元数据 get_max_partition,失败时降级用 MAX_PT 查询。
    """
    table_name = "dwd_video_element_contribution_analysis"
    try:
        table = client._odps.get_table(table_name)
        partition = table.get_max_partition()
        if partition is not None:
            name = getattr(partition, "name", "") or str(partition)
            for part in name.split(","):
                if part.startswith("dt="):
                    return part.split("=", 1)[1].strip("'\"")
    except Exception as e:
        logger.warning("[video_feature] get max partition by metadata failed:%s", e)

    rows = client.query(
        "SELECT MAX_PT('loghubods.dwd_video_element_contribution_analysis') AS dt",
        limit=1,
    )
    if rows and rows[0].get("dt"):
        return str(rows[0]["dt"])
    raise RuntimeError("无法获取 dwd_video_element_contribution_analysis 最新分区")


def _query_odps_rows(client, sql: str, columns: list[str]) -> list[dict]:
    """条件允许时不经 tunnel 读取 ODPS 行。

    部分运行环境无法访问配置的 tunnel endpoint；小批量特征查询使用实例读取器
    更可靠。
    """
    odps = getattr(client, "_odps", None)
    if odps is None:
        return client.query(sql)

    instance = odps.execute_sql(sql)
    instance.wait_for_success()
    rows: list[dict] = []
    with instance.open_reader(tunnel=False) as reader:
        for record in reader:
            rows.append({col: record[i] for i, col in enumerate(columns)})
    return rows


def _ensure_cache_table() -> None:
    """确保 MySQL 缓存表 video_element_feature_cache 存在且结构、索引符合预期(幂等迁移)。"""
    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(CREATE_CACHE_TABLE_SQL)
            cur.execute("SHOW COLUMNS FROM video_element_feature_cache LIKE 'element_dimension'")
            if not cur.fetchone():
                cur.execute(
                    """
ALTER TABLE video_element_feature_cache
ADD COLUMN element_dimension VARCHAR(50) NOT NULL DEFAULT '实质' COMMENT '元素维度'
AFTER standard_element
"""
                )
            cur.execute("SHOW INDEX FROM video_element_feature_cache WHERE Key_name = 'uk_dt_video_point_element'")
            index_rows = cur.fetchall()
            index_cols = [str(row.get("Column_name") or "") for row in index_rows]
            index_ok = (
                index_cols == ["dt", "video_id", "element_dimension", "point_type", "standard_element"]
                and str((index_rows[-1] or {}).get("Sub_part") or "") == "191"
            )
            if index_rows and not index_ok:
                cur.execute("ALTER TABLE video_element_feature_cache DROP INDEX uk_dt_video_point_element")
            cur.execute("SHOW COLUMNS FROM video_element_feature_cache LIKE 'standard_element'")
            standard_col = cur.fetchone()
            if standard_col and "varchar(1024)" not in str(standard_col.get("Type") or "").lower():
                cur.execute(
                    """
ALTER TABLE video_element_feature_cache
MODIFY COLUMN standard_element VARCHAR(1024) NOT NULL COMMENT '标准化元素或解构选题文本'
"""
                )
            cur.execute("SHOW COLUMNS FROM video_element_feature_cache LIKE 'is_miss'")
            if not cur.fetchone():
                cur.execute(
                    """
ALTER TABLE video_element_feature_cache
ADD COLUMN is_miss BOOLEAN NOT NULL DEFAULT FALSE COMMENT '是否为无特征负缓存'
AFTER contribution_score
"""
                )
            if not index_ok:
                cur.execute(
                    """
ALTER TABLE video_element_feature_cache
ADD UNIQUE KEY uk_dt_video_point_element
    (dt, video_id, element_dimension, point_type, standard_element(191))
"""
                )
    finally:
        conn.close()


def _read_cache(video_ids: list[int], dt: str) -> tuple[dict[int, list[VideoElementFeature]], set[int]]:
    """按 dt 从缓存表批量读特征。返回 (video_id→特征列表, 已缓存的 video_id 集合,含负缓存)。"""
    if not video_ids:
        return {}, set()
    from db.connection import get_connection

    out: dict[int, list[VideoElementFeature]] = {}
    cached_ids: set[int] = set()
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            for batch in _chunks(video_ids, 500):
                placeholders = ",".join(["%s"] * len(batch))
                cur.execute(
                    f"""
SELECT video_id, dt, element_dimension, point_type, standard_element, contribution_score, is_miss
FROM video_element_feature_cache
WHERE dt = %s
AND video_id IN ({placeholders})
ORDER BY video_id, contribution_score DESC
""",
                    [dt, *batch],
                )
                for row in cur.fetchall():
                    video_id = int(row["video_id"])
                    cached_ids.add(video_id)
                    if row.get("is_miss"):
                        continue
                    out.setdefault(video_id, []).append(VideoElementFeature(
                        video_id=video_id,
                        element_dimension=str(row["element_dimension"] or ""),
                        point_type=str(row["point_type"] or ""),
                        standard_element=str(row["standard_element"] or ""),
                        contribution_score=float(row["contribution_score"] or 0.0),
                        dt=str(row["dt"] or ""),
                    ))
    finally:
        conn.close()
    return out, cached_ids


def _write_cache(
    features: dict[int, list[VideoElementFeature]],
    miss_video_ids: Iterable[int],
    dt: str,
) -> None:
    """把 ODPS 查到的特征和无特征负缓存(is_miss=True)批量写入缓存表(upsert)。"""
    rows = [
        (
            feature.video_id,
            feature.dt,
            feature.point_type,
            feature.standard_element,
            feature.element_dimension,
            feature.contribution_score,
            False,
        )
        for feature_list in features.values()
        for feature in feature_list
    ]
    miss_rows = [
        (int(video_id), dt, "", "", "实质", 0.0, True)
        for video_id in miss_video_ids
    ]
    rows.extend(miss_rows)
    if not rows:
        return

    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.executemany(
                """
INSERT INTO video_element_feature_cache
    (video_id, dt, point_type, standard_element, element_dimension, contribution_score, is_miss)
VALUES (%s, %s, %s, %s, %s, %s, %s)
ON DUPLICATE KEY UPDATE
    contribution_score = VALUES(contribution_score),
    element_dimension = VALUES(element_dimension),
    is_miss = VALUES(is_miss),
    updated_at = CURRENT_TIMESTAMP
""",
                rows,
            )
    finally:
        conn.close()


def read_cached_video_element_features(
    video_ids: Iterable[int],
) -> dict[int, list[VideoElementFeature]]:
    """仅从本地数据库读取每个视频 ID 的最新缓存特征行。"""
    ids: list[int] = []
    seen: set[int] = set()
    for raw in video_ids:
        try:
            vid = int(raw)
        except (TypeError, ValueError):
            continue
        if vid in seen:
            continue
        seen.add(vid)
        ids.append(vid)
    if not ids:
        return {}

    _ensure_cache_table()
    from db.connection import get_connection

    out: dict[int, list[VideoElementFeature]] = {}
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            for batch in _chunks(ids, 500):
                placeholders = ",".join(["%s"] * len(batch))
                params = [*batch, *batch]
                cur.execute(
                    f"""
SELECT c.video_id, c.dt, c.element_dimension, c.point_type,
       c.standard_element, c.contribution_score, c.is_miss
FROM video_element_feature_cache c
JOIN (
    SELECT video_id, MAX(dt) AS dt
    FROM video_element_feature_cache
    WHERE video_id IN ({placeholders})
    GROUP BY video_id
) latest ON latest.video_id = c.video_id AND latest.dt = c.dt
WHERE c.video_id IN ({placeholders})
ORDER BY c.video_id, c.contribution_score DESC
""",
                    params,
                )
                for row in cur.fetchall():
                    if row.get("is_miss"):
                        continue
                    video_id = int(row["video_id"])
                    out.setdefault(video_id, []).append(VideoElementFeature(
                        video_id=video_id,
                        element_dimension=str(row["element_dimension"] or ""),
                        point_type=str(row["point_type"] or ""),
                        standard_element=str(row["standard_element"] or ""),
                        contribution_score=float(row["contribution_score"] or 0.0),
                        dt=str(row["dt"] or ""),
                    ))
    finally:
        conn.close()
    return out


def fetch_video_element_features(
    video_ids: Iterable[int],
    chunk_size: int = 100,
) -> dict[int, list[VideoElementFeature]]:
    """从 ODPS 最新分区读取各视频 ID 的召回特征。

    一个视频可以有多条可用元素，保留符合以下条件的全部记录：
    - 元素维度为“解构选题”；
    - 元素维度为“实质”，且点类型属于灵感点、关键点或目的点；
    - 贡献分不低于 0.8；
    - 标准化元素非空。
    """
    ids: list[int] = []
    seen: set[int] = set()
    for raw in video_ids:
        try:
            vid = int(raw)
        except (TypeError, ValueError):
            continue
        if vid in seen:
            continue
        seen.add(vid)
        ids.append(vid)

    if not ids:
        return {}

    client = _get_odps_client()
    if client is None:
        logger.warning("[video_feature] ODPS client unavailable, skip enrichment")
        return {}

    dt = _latest_partition_dt(client)
    _ensure_cache_table()
    out, cached_ids = _read_cache(ids, dt)
    missing_ids = [video_id for video_id in ids if video_id not in cached_ids]
    if not missing_ids:
        logger.info(
            "[video_feature] cache hit %d/%d videos dt=%s",
            len(cached_ids), len(ids), dt,
        )
        return out

    fetched: dict[int, list[VideoElementFeature]] = {}
    for batch in _chunks(missing_ids, max(1, chunk_size)):
        vid_list = ",".join(_quote_sql_string(str(v)) for v in batch)
        sql = f"""
SELECT  vid
        ,`元素维度` AS element_dimension
        ,`点类型` AS point_type
        ,`标准化元素` AS standard_element
        ,`解构选题` AS deconstruct_topic
        ,`贡献分` AS contribution_score
        ,dt
FROM    loghubods.dwd_video_element_contribution_analysis
WHERE   dt = {_quote_sql_string(dt)}
AND     vid IN ({vid_list})
AND     (
            (`元素维度` = '实质' AND `点类型` IN ('灵感点', '关键点', '目的点'))
            OR (`解构选题` IS NOT NULL AND `解构选题` <> '')
        )
AND     `贡献分` >= 0.8
ORDER BY vid, `贡献分` DESC
"""
        rows = _query_odps_rows(
            client,
            sql,
            ["vid", "element_dimension", "point_type", "standard_element", "deconstruct_topic", "contribution_score", "dt"],
        )
        seen_features: set[tuple[int, str, str, str]] = set()
        for row in rows:
            try:
                video_id = int(row.get("vid"))
            except (TypeError, ValueError):
                continue
            element_dimension = str(row.get("element_dimension") or "").strip()
            point_type = str(row.get("point_type") or "").strip()
            standard_element = str(row.get("standard_element") or "").strip()
            contribution_score = float(row.get("contribution_score") or 0.0)
            dt_value = str(row.get("dt") or "")

            if (
                element_dimension == "实质"
                and point_type in {"灵感点", "关键点", "目的点"}
                and standard_element
            ):
                key = (video_id, element_dimension, point_type, standard_element)
                if key not in seen_features:
                    seen_features.add(key)
                    fetched.setdefault(video_id, []).append(VideoElementFeature(
                        video_id=video_id,
                        element_dimension=element_dimension,
                        point_type=point_type,
                        standard_element=standard_element,
                        contribution_score=contribution_score,
                        dt=dt_value,
                    ))

            topic = str(row.get("deconstruct_topic") or "").strip()
            if topic:
                key = (video_id, "解构选题", "", topic)
                if key not in seen_features:
                    seen_features.add(key)
                    fetched.setdefault(video_id, []).append(VideoElementFeature(
                        video_id=video_id,
                        element_dimension="解构选题",
                        point_type="",
                        standard_element=topic,
                        contribution_score=contribution_score,
                        dt=dt_value,
                    ))

    fetched_ids = set(fetched.keys())
    miss_ids = [video_id for video_id in missing_ids if video_id not in fetched_ids]
    _write_cache(fetched, miss_ids, dt)
    out.update(fetched)
    logger.info(
        "[video_feature] dt=%s cache_hit=%d odps_hit=%d odps_miss=%d/%d element_rows=%d",
        dt, len(cached_ids), len(fetched), len(miss_ids), len(missing_ids),
        sum(len(v) for v in out.values()),
    )
    return out