import argparse import csv import hashlib import json import os import subprocess import time import urllib.request from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path YEAR = 2025 DATASET = "choucsan/Gaokao-Compass-11M" REVISION = "main" TREE_URL = f"https://huggingface.co/api/datasets/{DATASET}/tree/{REVISION}/data/{YEAR}?recursive=true&expand=false&limit=1000" LANDING_URL = f"https://huggingface.co/datasets/{DATASET}" USER_AGENT = "Mozilla/5.0 (compatible; ZhiyuantongDataAudit/1.0)" PROVINCES = { "anhui": "安徽省", "beijing": "北京市", "chongqing": "重庆市", "fujian": "福建省", "gansu": "甘肃省", "guangdong": "广东省", "guangxi": "广西壮族自治区", "guizhou": "贵州省", "hainan": "海南省", "hebei": "河北省", "heilongjiang": "黑龙江省", "henan": "河南省", "hubei": "湖北省", "hunan": "湖南省", "jiangsu": "江苏省", "jiangxi": "江西省", "jilin": "吉林省", "liaoning": "辽宁省", "neimenggu": "内蒙古自治区", "ningxia": "宁夏回族自治区", "qinghai": "青海省", "shaanxi": "陕西省", "shandong": "山东省", "shanghai": "上海市", "shanxi": "山西省", "sichuan": "四川省", "tianjin": "天津市", "xinjiang": "新疆维吾尔自治区", "xizang": "西藏自治区", "yunnan": "云南省", "zhejiang": "浙江省", } EXPECTED_HEADERS = { "enrollment-plan.csv": [ "province", "year", "batch", "category", "university_code", "university_name", "major_group", "major_code", "major_name", "major_note", "subject_req", "plan_count", "duration", "tuition", ], "school-admission.csv": [ "province", "year", "category", "batch", "university_code", "university_name", "major_group", "subject_req", "min_score", "min_rank", "control_score", "score_diff", "admit_count", "school_province", "school_nature", "is_985", "is_211", ], "major-admission.csv": [ "province", "year", "category", "batch", "university_code", "university_name", "major_group", "major_code", "major_name", "major_note", "subject_req", "min_score", "min_rank", "max_score", "avg_score", "admit_count", "school_province", "school_nature", "is_985", "is_211", ], "score-range.csv": [ "province", "year", "category", "batch", "control_score", "score", "score_range", "segment_count", "cumulative_count", "rank_range", ], } HEADER_VARIANTS = { ("fujian", "enrollment-plan.csv"): [ "province", "year", "batch", "category", "university_code", "university_name", "major_group", "major_group", "major_code", "major_name", "major_note", "subject_req", "plan_count", "duration", "tuition", ], ("shanghai", "enrollment-plan.csv"): [ "province", "year", "batch", "category", "university_code", "university_name", "major_group", "major_group", "major_code", "major_name", "major_note", "subject_req", "plan_count", "duration", "tuition", ], ("guangdong", "enrollment-plan.csv"): [ "province", "year", "batch", "category", "university_code", "university_name", "major_group", "major_code", "major_name", "major_note", "major_note", "subject_req", "plan_count", "duration", "tuition", ], } STAGING = { "enrollment-plan.csv": ( "stg_gc_enrollment_plans", "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw," "group_code,major_code,major_name,major_note,subject_requirements,plan_count_text,study_duration,tuition_text", ), "school-admission.csv": ( "stg_gc_school_admissions", "@raw_province,@raw_year,subject_category,batch_name,source_school_code,school_name_raw," "group_code,subject_requirements,min_score_text,min_rank_text,control_score_text,score_diff_text," "admit_count_text,school_province,school_nature,is_985_text,is_211_text", ), "major-admission.csv": ( "stg_gc_major_admissions", "@raw_province,@raw_year,subject_category,batch_name,source_school_code,school_name_raw," "group_code,major_code,major_name,major_note,subject_requirements,min_score_text,min_rank_text," "max_score_text,avg_score_text,admit_count_text,school_province,school_nature,is_985_text,is_211_text", ), "score-range.csv": ( "stg_gc_score_ranges", "@raw_province,@raw_year,subject_category,batch_name,control_score_text,score_text,score_range_text," "segment_count_text,cumulative_count_text,rank_range_text", ), } def sql_string(value): if value is None: return "NULL" text = str(value).replace("\\", "\\\\").replace("'", "\\'") text = text.replace("\x00", "").replace("\r", " ").replace("\n", " ") return f"'{text}'" def request_bytes(url, timeout=180): request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT}) last_error = None for attempt in range(3): try: with urllib.request.urlopen(request, timeout=timeout) as response: return response.read() except Exception as error: # noqa: BLE001 - 重试后由调用方报告完整URL last_error = error if attempt < 2: time.sleep(2 ** attempt) raise RuntimeError(f"下载失败: {url}: {last_error}") def sha256(path): digest = hashlib.sha256() with path.open("rb") as stream: for block in iter(lambda: stream.read(1024 * 1024), b""): digest.update(block) return digest.hexdigest() def download_one(item, raw_root, refresh): relative = Path(item["path"]).relative_to(f"data/{YEAR}") target = raw_root / relative target.parent.mkdir(parents=True, exist_ok=True) expected_size = int(item["size"]) if not refresh and target.exists() and target.stat().st_size == expected_size: return target url = f"https://huggingface.co/datasets/{DATASET}/resolve/{REVISION}/{item['path']}" content = request_bytes(url) if len(content) != expected_size: raise RuntimeError(f"文件长度异常: {item['path']}，期望{expected_size}，实际{len(content)}") partial = target.with_suffix(target.suffix + ".part") partial.write_bytes(content) os.replace(partial, target) return target def csv_summary(path): expected = HEADER_VARIANTS.get((path.parent.name, path.name), EXPECTED_HEADERS[path.name]) row_count = 0 invalid_width = 0 with path.open("r", encoding="utf-8-sig", newline="") as stream: reader = csv.reader(stream) header = next(reader, None) if header != expected: raise RuntimeError(f"字段不匹配: {path}\n期望: {expected}\n实际: {header}") for row in reader: row_count += 1 if len(row) != len(expected): invalid_width += 1 if invalid_width: raise RuntimeError(f"CSV列数异常: {path}，异常行数{invalid_width}") return row_count def load_mapping(path): slug = path.parent.name if path.name == "enrollment-plan.csv" and slug in {"fujian", "shanghai"}: return ( "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw," "@composite_group_code,group_code,major_code,major_name,major_note,subject_requirements," "plan_count_text,study_duration,tuition_text", "", ) if path.name == "enrollment-plan.csv" and slug == "guangdong": return ( "@raw_province,@raw_year,batch_name,subject_category,source_school_code,school_name_raw," "group_code,major_code,major_name,@major_note_one,@major_note_two,subject_requirements," "plan_count_text,study_duration,tuition_text", ", major_note=NULLIF(CONCAT_WS('；', NULLIF(TRIM(@major_note_one), ''), NULLIF(TRIM(@major_note_two), '')), '')", ) return STAGING[path.name][1], "" def mysql_command(args, extra=None): if not os.environ.get("MYSQL_PWD"): raise RuntimeError("必须通过MYSQL_PWD环境变量提供数据库密码") command = [ args.mysql_bin, "--protocol=TCP", "-h", args.host, "-P", str(args.port), "-u", args.user, "--default-character-set=utf8mb4", ] if extra: command.extend(extra) return command def mysql_run(args, sql, local_infile=False, capture=True): extra = ["--local-infile=1"] if local_infile else [] process = subprocess.run( mysql_command(args, extra), input=sql, text=True, encoding="utf-8", capture_output=capture, ) if process.returncode: raise RuntimeError(process.stderr.strip() or process.stdout.strip()) return process.stdout def mysql_run_file(args, path): mysql_run(args, path.read_text(encoding="utf-8")) def build_load_sql(csv_files, summaries, database): statements = [ "SET NAMES utf8mb4;", f"USE `{database}`;", "SET SESSION sql_log_bin=0;", "SET SESSION sql_mode='STRICT_TRANS_TABLES,ERROR_FOR_DIVISION_BY_ZERO,NO_ENGINE_SUBSTITUTION';", "TRUNCATE TABLE stg_gc_enrollment_plans;", "TRUNCATE TABLE stg_gc_school_admissions;", "TRUNCATE TABLE stg_gc_major_admissions;", "TRUNCATE TABLE stg_gc_score_ranges;", ] for path in csv_files: slug = path.parent.name province = PROVINCES[slug] filename = path.name source_url = f"https://huggingface.co/datasets/{DATASET}/resolve/{REVISION}/data/{YEAR}/{slug}/{filename}" digest = sha256(path) title_map = { "enrollment-plan.csv": "招生计划", "school-admission.csv": "院校投档线", "major-admission.csv": "专业录取线", "score-range.csv": "一分一段", } notes = ( f"GaokaoCompass开放数据集{YEAR}年{province}{title_map[filename]}CSV；" "来源方说明其由各省教育考试院、阳光高考等公开渠道清洗整理。" "本系统将该文件作为待审核快照，不替代省级考试院原始公告。" ) statements.append(f""" INSERT INTO data_sources ( source_type, publisher, title, landing_page_url, file_url, published_on, retrieved_at, sha256, raw_file_path, notes ) VALUES ( 'open_dataset_csv', 'GaokaoCompass', {sql_string(f'{province}{YEAR}年{title_map[filename]}')}, {sql_string(LANDING_URL)}, {sql_string(source_url)}, NULL, NOW(), {sql_string(digest)}, {sql_string(str(path.resolve()))}, {sql_string(notes)} ) AS incoming ON DUPLICATE KEY UPDATE retrieved_at=incoming.retrieved_at, raw_file_path=incoming.raw_file_path, file_url=incoming.file_url, notes=incoming.notes; SET @source_id=(SELECT source_id FROM data_sources WHERE sha256={sql_string(digest)}); SET @row_no=0; """) table = STAGING[filename][0] columns, extra_set = load_mapping(path) local_path = path.resolve().as_posix() statements.append(f""" LOAD DATA LOCAL INFILE {sql_string(local_path)} INTO TABLE {table} CHARACTER SET utf8mb4 FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '"' ESCAPED BY '"' LINES TERMINATED BY '\n' IGNORE 1 LINES ({columns}) SET source_id=@source_id, source_row_number=(@row_no:=@row_no+1), province_name={sql_string(province)}, admission_year=CAST(TRIM(@raw_year) AS UNSIGNED){extra_set}; """) statements.append("SELECT 'staging_loaded' AS status;") return "\n".join(statements) def main(): parser = argparse.ArgumentParser(description="下载并覆盖导入GaokaoCompass 2025全国招录数据") parser.add_argument("--mysql-bin", required=True) parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=3306) parser.add_argument("--user", default="root") parser.add_argument("--database", default="zhiyuantong") parser.add_argument("--root", default=str(Path(__file__).resolve().parents[1])) parser.add_argument("--refresh", action="store_true", help="强制重新下载文件") parser.add_argument("--workers", type=int, default=8) parser.add_argument("--province", choices=sorted(PROVINCES), help="只覆盖指定省份目录（英文目录名）") args = parser.parse_args() root = Path(args.root).resolve() raw_root = root / "data" / "raw" / "gaokao_compass" / str(YEAR) raw_root.mkdir(parents=True, exist_ok=True) tree = json.loads(request_bytes(TREE_URL, timeout=60).decode("utf-8")) items = [ item for item in tree if item.get("type") == "file" and Path(item["path"]).name in set(EXPECTED_HEADERS) | {"meta.json"} ] expected_provinces = {Path(item["path"]).parts[-2] for item in items if Path(item["path"]).name == "enrollment-plan.csv"} if expected_provinces != set(PROVINCES): raise RuntimeError(f"招生计划省份不完整: {sorted(set(PROVINCES) - expected_provinces)}") print(f"开始下载/校验{len(items)}个文件，共{sum(int(item['size']) for item in items):,}字节") downloaded = [] with ThreadPoolExecutor(max_workers=max(1, min(args.workers, 12))) as executor: futures = {executor.submit(download_one, item, raw_root, args.refresh): item for item in items} for index, future in enumerate(as_completed(futures), 1): downloaded.append(future.result()) if index % 15 == 0 or index == len(items): print(f"已完成 {index}/{len(items)} 个文件") csv_files = sorted(path for path in downloaded if path.suffix.lower() == ".csv") summaries = {str(path): csv_summary(path) for path in csv_files} manifest = { "dataset": DATASET, "revision": REVISION, "year": YEAR, "landing_url": LANDING_URL, "downloaded_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "files": [ { "path": str(path.relative_to(root)).replace("\\", "/"), "size": path.stat().st_size, "sha256": sha256(path), "rows": summaries.get(str(path)), } for path in sorted(downloaded) ], } (raw_root / "manifest.json").write_text( json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8" ) print(f"CSV结构校验通过：{len(csv_files)}个文件，共{sum(summaries.values()):,}行") migration = root / "database" / "migrations" / "007_add_gaokao_compass_2025_tables.sql" replacement = root / "database" / "import_gaokao_compass_2025.sql" migration_state = mysql_run(args, f""" SELECT COUNT(*) FROM information_schema.columns WHERE table_schema={sql_string(args.database)} AND table_name='admission_groups' AND column_name='batch_name'; """).strip().splitlines()[-1] if migration_state == "0": mysql_run_file(args, migration) else: print("数据库007迁移已存在，跳过重复执行") mysql_run(args, "SET GLOBAL local_infile=ON;") try: province_slugs = [args.province] if args.province else sorted(PROVINCES) for index, slug in enumerate(province_slugs, 1): province_files = [path for path in csv_files if path.parent.name == slug] if not any(path.name == "enrollment-plan.csv" for path in province_files): raise RuntimeError(f"{PROVINCES[slug]}缺少招生计划文件，不能执行省级覆盖") load_output = mysql_run( args, build_load_sql(province_files, summaries, args.database), local_infile=True ) if "staging_loaded" not in load_output: raise RuntimeError(f"{PROVINCES[slug]}暂存表装载未返回完成标志") stage_counts = mysql_run(args, f""" USE `{args.database}`; SELECT 'enrollment-plan.csv', COUNT(*) FROM stg_gc_enrollment_plans UNION ALL SELECT 'school-admission.csv', COUNT(*) FROM stg_gc_school_admissions UNION ALL SELECT 'major-admission.csv', COUNT(*) FROM stg_gc_major_admissions UNION ALL SELECT 'score-range.csv', COUNT(*) FROM stg_gc_score_ranges; """) actual = {} for line in stage_counts.splitlines(): parts = line.split("\t") if len(parts) == 2 and parts[1].isdigit(): actual[parts[0]] = int(parts[1]) expected_by_type = { filename: sum( summaries[str(path)] for path in province_files if path.name == filename ) for filename in EXPECTED_HEADERS } if actual != expected_by_type: raise RuntimeError( f"{PROVINCES[slug]}暂存表行数校验失败，期望{expected_by_type}，实际{actual}" ) distinct_provinces = mysql_run(args, f""" USE `{args.database}`; SELECT COUNT(DISTINCT province_name), MIN(province_name) FROM stg_gc_enrollment_plans; """).strip().splitlines()[-1].split("\t") if distinct_provinces != ["1", PROVINCES[slug]]: raise RuntimeError(f"暂存表省份边界校验失败：{distinct_provinces}") mysql_run_file(args, replacement) print( f"[{index}/{len(province_slugs)}] {PROVINCES[slug]}覆盖完成：" + ", ".join(f"{name}={count:,}" for name, count in actual.items()), flush=True, ) finally: mysql_run(args, "SET GLOBAL local_infile=OFF;") verification = mysql_run(args, f""" USE `{args.database}`; SELECT '专业招生计划', COUNT(*) FROM admission_major_plans plans JOIN admission_groups groups_ ON groups_.admission_group_id=plans.admission_group_id WHERE groups_.admission_year=2025 UNION ALL SELECT '院校专业组', COUNT(*) FROM admission_groups WHERE admission_year=2025 UNION ALL SELECT '院校投档线', COUNT(*) FROM admission_cutoff_lines WHERE admission_year=2025 UNION ALL SELECT '专业录取线', COUNT(*) FROM admission_major_lines WHERE admission_year=2025 UNION ALL SELECT '一分一段', COUNT(*) FROM score_rank_points WHERE admission_year=2025 UNION ALL SELECT '省级控制线', COUNT(*) FROM control_score_lines WHERE admission_year=2025; SELECT '覆盖省份', COUNT(DISTINCT province_name) FROM admission_major_plans plans JOIN admission_groups groups_ ON groups_.admission_group_id=plans.admission_group_id WHERE groups_.admission_year=2025; """) print("覆盖导入完成：") print(verification.strip()) if __name__ == "__main__": main()