Commit 98b4f92e 98b4f92e3cf7d503e7e4635d9a9ab39aaa4bce58 by 沈秋雨

perf(etl): 优化批量更新与批处理大小

- 将默认批处理大小 BATCH_SIZE 从 1 增大到 1000 提高效率
- 添加 BACKFILL_BATCH_SIZE 环境变量支持,默认为 5000 用于补数据批量控制
- 使用 SQL 的 VALUES 语法替换 executemany 实现批量更新 yinyan_song_records 平台字段
- 修改 runner 中补数据查询改用 BACKFILL_BATCH_SIZE 以提升性能
- 更新相关单元测试以适配新的批量更新 SQL 语句和参数格式
1 parent d8728ac1
......@@ -43,4 +43,5 @@ PLATFORM_KUGOU = '2'
PLATFORM_NETEASE = '4'
PLATFORMS = [PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE]
BATCH_SIZE = 1
BATCH_SIZE = 1000
BACKFILL_BATCH_SIZE = int(os.environ.get('BACKFILL_BATCH_SIZE', '5000'))
......
......@@ -3,7 +3,7 @@ import json
import logging
from tqdm import tqdm
from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, OSS_CONFIG
from .config import PLATFORM_QQ, PLATFORM_KUGOU, PLATFORM_NETEASE, BATCH_SIZE, BACKFILL_BATCH_SIZE, OSS_CONFIG
from .connections import get_hk_songs_conn, get_source_conn, get_spider_conn, get_pg_conn, get_oss_bucket
from .reader import (
iter_hk_songs_batches,
......@@ -548,7 +548,7 @@ def backfill_yinyan_record_platforms(max_batches: int | None = None) -> None:
pbar = tqdm(desc='backfill-yinyan-platforms')
while max_batches is None or batch_index < max_batches:
with pg_conn.cursor() as pg_cur:
rows = fetch_yinyan_records_missing_platform(pg_cur, BATCH_SIZE)
rows = fetch_yinyan_records_missing_platform(pg_cur, BACKFILL_BATCH_SIZE)
if not rows:
break
......
......@@ -56,15 +56,20 @@ def fetch_yinyan_records_missing_platform(cur, limit: int) -> list[dict]:
def update_yinyan_record_platforms(cur, records: list[dict]) -> None:
if not records:
return
cur.executemany(
"""
UPDATE yinyan_song_records
SET platform = %s
WHERE song_id = %s
AND record_id = %s
AND platform IS NULL
values_sql = ', '.join(['(%s::bigint, %s::bigint, %s::varchar)'] * len(records))
params = []
for r in records:
params.extend([r['song_id'], r['record_id'], r['platform']])
cur.execute(
f"""
UPDATE yinyan_song_records AS ysr
SET platform = v.platform
FROM (VALUES {values_sql}) AS v(song_id, record_id, platform)
WHERE ysr.song_id = v.song_id
AND ysr.record_id = v.record_id
AND ysr.platform IS NULL
""",
[(r['platform'], r['song_id'], r['record_id']) for r in records],
tuple(params),
)
......
......@@ -97,10 +97,12 @@ def test_update_yinyan_record_platforms_fills_only_null_platform_rows():
{'song_id': 11, 'record_id': 101, 'platform': '2'},
])
sql, rows = cur.executemany.call_args[0]
assert 'SET platform = %s' in sql
assert 'AND platform IS NULL' in sql
assert rows == [('1', 10, 100), ('2', 11, 101)]
sql, params = cur.execute.call_args[0]
assert 'UPDATE yinyan_song_records AS ysr' in sql
assert 'FROM (VALUES (%s::bigint, %s::bigint, %s::varchar), (%s::bigint, %s::bigint, %s::varchar))' in sql
assert 'SET platform = v.platform' in sql
assert 'AND ysr.platform IS NULL' in sql
assert params == (10, 100, '1', 11, 101, '2')
def test_upsert_netease_songs_writes_album_json_column():
......