|
|
@@ -0,0 +1,175 @@
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+from sqlalchemy import case, func, select
|
|
|
+from sqlalchemy.dialects.mysql import insert
|
|
|
+from sqlalchemy.orm import Session
|
|
|
+
|
|
|
+from supply_infra.db.models.channel_content_data import ChannelContentData
|
|
|
+from supply_infra.db.models.global_category_content_weight import (
|
|
|
+ GlobalCategoryContentWeight,
|
|
|
+)
|
|
|
+from supply_infra.db.models.global_v2 import GlobalCategoryV2, GlobalSourceElementData
|
|
|
+
|
|
|
+_INSERT_BATCH_SIZE = 500
|
|
|
+
|
|
|
+
|
|
|
+class GlobalCategoryContentWeightRepository:
|
|
|
+ def __init__(self, session: Session) -> None:
|
|
|
+ self.session = session
|
|
|
+
|
|
|
+ def count_by_biz_dt(self, biz_dt: str) -> int:
|
|
|
+ stmt = select(func.count()).where(GlobalCategoryContentWeight.biz_dt == biz_dt)
|
|
|
+ return int(self.session.scalar(stmt) or 0)
|
|
|
+
|
|
|
+ def get_latest_completed_biz_dt(self) -> str | None:
|
|
|
+ """Return the newest date whose node rows are all completed."""
|
|
|
+ completed_count = func.sum(
|
|
|
+ case((GlobalCategoryContentWeight.status == "completed", 1), else_=0)
|
|
|
+ )
|
|
|
+ stmt = (
|
|
|
+ select(GlobalCategoryContentWeight.biz_dt)
|
|
|
+ .group_by(GlobalCategoryContentWeight.biz_dt)
|
|
|
+ .having(completed_count == func.count())
|
|
|
+ .order_by(GlobalCategoryContentWeight.biz_dt.desc())
|
|
|
+ .limit(1)
|
|
|
+ )
|
|
|
+ value = self.session.scalar(stmt)
|
|
|
+ return str(value) if value else None
|
|
|
+
|
|
|
+ def list_completed_models_by_biz_dt(self, biz_dt: str) -> list[GlobalCategoryContentWeight]:
|
|
|
+ return list(
|
|
|
+ self.session.scalars(
|
|
|
+ select(GlobalCategoryContentWeight).where(
|
|
|
+ GlobalCategoryContentWeight.biz_dt == biz_dt,
|
|
|
+ GlobalCategoryContentWeight.status == "completed",
|
|
|
+ )
|
|
|
+ ).all()
|
|
|
+ )
|
|
|
+
|
|
|
+ def direct_stats(self, biz_dt: str) -> dict[int, dict[str, float | int]]:
|
|
|
+ contribution = GlobalSourceElementData.contribution
|
|
|
+ stmt = (
|
|
|
+ select(
|
|
|
+ GlobalSourceElementData.global_category_stable_id.label("stable_id"),
|
|
|
+ func.count().label("source_element_count"),
|
|
|
+ func.sum(contribution).label("contribution_sum"),
|
|
|
+ func.sum(ChannelContentData.avg_read_rate * contribution).label(
|
|
|
+ "avg_read_rate_weighted_sum"
|
|
|
+ ),
|
|
|
+ func.sum(ChannelContentData.like_rate * contribution).label(
|
|
|
+ "like_rate_weighted_sum"
|
|
|
+ ),
|
|
|
+ func.sum(ChannelContentData.read_rate * contribution).label(
|
|
|
+ "read_rate_weighted_sum"
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ .join(
|
|
|
+ ChannelContentData,
|
|
|
+ (ChannelContentData.channel_content_id == GlobalSourceElementData.post_id)
|
|
|
+ & (ChannelContentData.dt == biz_dt),
|
|
|
+ )
|
|
|
+ .where(
|
|
|
+ GlobalSourceElementData.contribution.is_not(None),
|
|
|
+ GlobalSourceElementData.global_category_stable_id.is_not(None),
|
|
|
+ )
|
|
|
+ .group_by(GlobalSourceElementData.global_category_stable_id)
|
|
|
+ )
|
|
|
+ return {
|
|
|
+ int(row.stable_id): {
|
|
|
+ "source_element_count": int(row.source_element_count or 0),
|
|
|
+ "contribution_sum": float(row.contribution_sum or 0.0),
|
|
|
+ "avg_read_rate_weighted_sum": float(row.avg_read_rate_weighted_sum or 0.0),
|
|
|
+ "like_rate_weighted_sum": float(row.like_rate_weighted_sum or 0.0),
|
|
|
+ "read_rate_weighted_sum": float(row.read_rate_weighted_sum or 0.0),
|
|
|
+ }
|
|
|
+ for row in self.session.execute(stmt).all()
|
|
|
+ }
|
|
|
+
|
|
|
+ def initialize_day(self, biz_dt: str) -> dict[str, int]:
|
|
|
+ existing = self.count_by_biz_dt(biz_dt)
|
|
|
+ if existing:
|
|
|
+ return {"initialized": 0, "existing": existing}
|
|
|
+
|
|
|
+ categories = self.session.execute(
|
|
|
+ select(
|
|
|
+ GlobalCategoryV2.stable_id,
|
|
|
+ GlobalCategoryV2.parent_stable_id,
|
|
|
+ GlobalCategoryV2.level,
|
|
|
+ )
|
|
|
+ ).all()
|
|
|
+ direct_by_id = self.direct_stats(biz_dt)
|
|
|
+ rows: list[dict[str, Any]] = []
|
|
|
+ for category in categories:
|
|
|
+ stable_id = int(category.stable_id)
|
|
|
+ direct = direct_by_id.get(stable_id, {})
|
|
|
+ rows.append(
|
|
|
+ {
|
|
|
+ "stable_id": stable_id,
|
|
|
+ "parent_stable_id": (
|
|
|
+ int(category.parent_stable_id)
|
|
|
+ if category.parent_stable_id not in (None, 0)
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ "level": int(category.level) if category.level is not None else None,
|
|
|
+ "biz_dt": biz_dt,
|
|
|
+ "status": "pending",
|
|
|
+ "direct_source_element_count": int(direct.get("source_element_count", 0)),
|
|
|
+ "direct_contribution_sum": float(direct.get("contribution_sum", 0.0)),
|
|
|
+ "direct_avg_read_rate_weighted_sum": float(
|
|
|
+ direct.get("avg_read_rate_weighted_sum", 0.0)
|
|
|
+ ),
|
|
|
+ "direct_like_rate_weighted_sum": float(
|
|
|
+ direct.get("like_rate_weighted_sum", 0.0)
|
|
|
+ ),
|
|
|
+ "direct_read_rate_weighted_sum": float(
|
|
|
+ direct.get("read_rate_weighted_sum", 0.0)
|
|
|
+ ),
|
|
|
+ }
|
|
|
+ )
|
|
|
+
|
|
|
+ inserted = 0
|
|
|
+ for start in range(0, len(rows), _INSERT_BATCH_SIZE):
|
|
|
+ batch = rows[start : start + _INSERT_BATCH_SIZE]
|
|
|
+ result = self.session.execute(
|
|
|
+ insert(GlobalCategoryContentWeight).values(batch).prefix_with("IGNORE")
|
|
|
+ )
|
|
|
+ inserted += int(result.rowcount or 0)
|
|
|
+ return {"initialized": inserted, "existing": 0}
|
|
|
+
|
|
|
+ def list_day(self, biz_dt: str) -> list[dict[str, Any]]:
|
|
|
+ rows = self.session.scalars(
|
|
|
+ select(GlobalCategoryContentWeight)
|
|
|
+ .where(GlobalCategoryContentWeight.biz_dt == biz_dt)
|
|
|
+ .order_by(
|
|
|
+ GlobalCategoryContentWeight.level.desc(),
|
|
|
+ GlobalCategoryContentWeight.stable_id,
|
|
|
+ )
|
|
|
+ ).all()
|
|
|
+ fields = (
|
|
|
+ "id",
|
|
|
+ "stable_id",
|
|
|
+ "parent_stable_id",
|
|
|
+ "level",
|
|
|
+ "status",
|
|
|
+ "attempt_count",
|
|
|
+ "direct_source_element_count",
|
|
|
+ "source_element_count",
|
|
|
+ "direct_contribution_sum",
|
|
|
+ "contribution_sum",
|
|
|
+ "direct_avg_read_rate_weighted_sum",
|
|
|
+ "avg_read_rate_weighted_sum",
|
|
|
+ "avg_read_rate_score",
|
|
|
+ "direct_like_rate_weighted_sum",
|
|
|
+ "like_rate_weighted_sum",
|
|
|
+ "like_rate_score",
|
|
|
+ "direct_read_rate_weighted_sum",
|
|
|
+ "read_rate_weighted_sum",
|
|
|
+ "read_rate_score",
|
|
|
+ )
|
|
|
+ return [{field: getattr(row, field) for field in fields} for row in rows]
|
|
|
+
|
|
|
+ def save_completed(self, rows: list[dict[str, Any]]) -> None:
|
|
|
+ if rows:
|
|
|
+ self.session.bulk_update_mappings(GlobalCategoryContentWeight, rows)
|