如何高效实现从大型PGN数据库提取随机大师棋局的函数？

你需要实现函数从百万级的caissabase国际象棋PGN数据库中提取随机大师棋局，当前使用的两种方法存在明显问题：

• 水塘抽样法（extract_random_game）：遍历全部棋局耗时极久，且容易因单个坏棋局导致程序崩溃，调试困难。
• 按索引提取法（extract_game_by_index）：小索引提取正常，但提取大索引（如第1000000局）时需要从头遍历，耗时过长。

1. 预生成棋局位置索引文件（推荐） #

核心思路是先遍历一次PGN文件，记录每个棋局在文件中的起始偏移量，保存为索引文件。后续提取时直接通过偏移量跳转读取，无需从头遍历，大幅提升速度。

生成索引文件 #

import os
import chess.pgn
import pickle

def build_pgn_index(pgn_file, index_file):
    index = []
    with open(pgn_file, 'rb') as f:
        offset = 0
        while True:
            try:
                game = chess.pgn.read_game(f)
                if game is None:
                    break
                index.append(offset)
                offset = f.tell()
            except Exception as e:
                print(f"跳过损坏的棋局，错误信息: {e}")
                offset = f.tell()
                continue
    with open(index_file, 'wb') as idx_f:
        pickle.dump(index, idx_f)
    print(f"索引生成完成，共记录{len(index)}个棋局")

提取随机棋局 #

import random
import pickle
import chess.pgn

def extract_random_game_with_index(pgn_file, index_file, output_file):
    with open(index_file, 'rb') as idx_f:
        index = pickle.load(idx_f)
    
    if not index:
        print("索引文件中无有效棋局记录")
        return
    
    random_offset = random.choice(index)
    with open(pgn_file, 'rb') as f:
        f.seek(random_offset)
        try:
            game = chess.pgn.read_game(f)
            if game:
                with open(output_file, 'w') as out_f:
                    out_f.write(str(game))
                print(f"随机棋局已保存至: {output_file}")
            else:
                print("读取选中的棋局失败")
        except Exception as e:
            print(f"读取棋局出错: {e}")

按索引提取指定棋局 #

import pickle
import chess.pgn

def extract_game_by_index_with_index(pgn_file, index_file, game_index, output_file):
    with open(index_file, 'rb') as idx_f:
        index = pickle.load(idx_f)
    
    total_games = len(index)
    if game_index < 1 or game_index > total_games:
        print(f"索引超出范围，当前数据库共有{total_games}个棋局")
        return
    
    target_offset = index[game_index - 1]  # 索引文件为0-based，用户输入为1-based
    with open(pgn_file, 'rb') as f:
        f.seek(target_offset)
        try:
            game = chess.pgn.read_game(f)
            if game:
                with open(output_file, 'w') as out_f:
                    out_f.write(str(game))
                print(f"第{game_index}局已保存至: {output_file}")
            else:
                print("读取选中的棋局失败")
        except Exception as e:
            print(f"读取棋局出错: {e}")

2. 优化流式读取的稳定性 #

如果不想生成额外索引文件，可以通过异常捕获避免程序崩溃，同时优化读取逻辑：

import random
import chess.pgn

def extract_random_game_safe(pgn_file, output_file):
    random_game = None
    num_games = 0
    with open(pgn_file) as pgn:
        while True:
            try:
                game = chess.pgn.read_game(pgn)
                if game is None:
                    break
                num_games += 1
                if random.randint(1, num_games) == 1:
                    random_game = game
            except Exception as e:
                print(f"跳过损坏的棋局: {e}")
                continue
    
    if num_games == 0:
        print("PGN文件中未找到任何棋局")
        return
    
    with open(output_file, 'w') as new_pgn:
        new_pgn.write(str(random_game))
    print("随机棋局已保存至:", output_file)

3. 导入数据库管理（适合频繁查询场景） #

将PGN导入SQLite或PostgreSQL，利用数据库的索引和查询能力快速获取棋局：

导入PGN到SQLite #

import sqlite3
import chess.pgn

def import_pgn_to_sqlite(pgn_file, db_file):
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    cursor.execute('''CREATE TABLE IF NOT EXISTS games
                      (id INTEGER PRIMARY KEY AUTOINCREMENT, pgn_text TEXT)''')
    
    with open(pgn_file) as pgn:
        count = 0
        while True:
            try:
                game = chess.pgn.read_game(pgn)
                if game is None:
                    break
                cursor.execute("INSERT INTO games (pgn_text) VALUES (?)", (str(game),))
                count += 1
                if count % 10000 == 0:
                    conn.commit()
                    print(f"已导入{count}个棋局")
            except Exception as e:
                print(f"跳过损坏的棋局: {e}")
                continue
        conn.commit()
    print(f"导入完成，共{count}个棋局")
    conn.close()

从数据库提取随机棋局 #

import sqlite3

def extract_random_game_from_db(db_file, output_file):
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    cursor.execute("SELECT pgn_text FROM games ORDER BY RANDOM() LIMIT 1")
    result = cursor.fetchone()
    
    if result:
        with open(output_file, 'w') as out_f:
            out_f.write(result[0])
        print(f"随机棋局已保存至: {output_file}")
    else:
        print("数据库中无有效棋局")
    conn.close()

方案对比 #

内容的提问来源于stack exchange，提问作者Curtis