速いSQLは安いSQLなのか？ TiDB Cloudで検証してみた

"""TiDB Cloud Starter で、SQL の実行時間と RU を対にして計測する。

接続情報は環境変数 TIDB_HOST, TIDB_PORT, TIDB_USER, TIDB_PASSWORD から読む。

    python ru_bench.py setup             # テーブル作成とデータ投入
    python ru_bench.py stats             # 統計情報の反映を確認 (stats:pseudo が False なら反映済み)
    python ru_bench.py bench ab|conc|lock  # 計測
"""

import json
import os
import ssl
import statistics
import sys
import threading
import time

import pymysql

DB = "ru_lab"

SETUP = [
    f"CREATE DATABASE IF NOT EXISTS {DB}",
    f"USE {DB}",
    "CREATE TABLE IF NOT EXISTS digit (n INT PRIMARY KEY)",
    "INSERT IGNORE INTO digit VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)",
    """CREATE TABLE IF NOT EXISTS orders (
      id BIGINT PRIMARY KEY,
      customer_id INT NOT NULL,
      status TINYINT NOT NULL,
      ordered_at DATETIME NOT NULL,
      note VARCHAR(128) NOT NULL,
      KEY idx_history (customer_id, status, ordered_at DESC, id DESC)
    )""",
    """CREATE TABLE IF NOT EXISTS order_items (
      order_id BIGINT NOT NULL,
      line_no TINYINT NOT NULL,
      qty TINYINT NOT NULL,
      unit_price INT NOT NULL,
      PRIMARY KEY (order_id, line_no)
    )""",
]

# 注文 10,000 件 (うち 1〜5,000 は顧客 42 の支払い済み注文) と、注文ごとに 10 件の明細
LOAD = [
    """INSERT INTO orders
    SELECT x.id,
           IF(x.id <= 5000, 42, IF(MOD(x.id, 100) = 42, 43, MOD(x.id, 100))),
           IF(x.id <= 5000, 1, IF(MOD(x.id, 7) = 0, 0, 1)),
           DATE_ADD('2026-09-01', INTERVAL MOD(x.id, 30 * 24 * 60) MINUTE),
           RPAD(CONCAT('memo-', x.id), 128, 'x')
    FROM (
      SELECT 1 + a.n + 10*b.n + 100*c.n + 1000*d.n AS id
      FROM digit a CROSS JOIN digit b CROSS JOIN digit c CROSS JOIN digit d
    ) AS x""",
    """INSERT INTO order_items
    SELECT o.id, e.n + 1, 1 + MOD(e.n, 3), 100 + MOD(o.id + e.n, 500)
    FROM orders AS o CROSS JOIN digit AS e""",
]

# 顧客 42 の支払い済み注文を新しい順に 100 件、合計金額つきで返す
QUERY_A = """SELECT o.id, o.ordered_at, o.note,
       COALESCE(SUM(i.qty * i.unit_price), 0) AS total_yen
FROM orders AS o
LEFT JOIN order_items AS i ON i.order_id = o.id
WHERE o.customer_id = 42 AND o.status = 1
GROUP BY o.id, o.ordered_at, o.note
ORDER BY o.ordered_at DESC, o.id DESC
LIMIT 100"""

QUERY_B = """SELECT p.id, p.ordered_at, p.note,
       COALESCE(SUM(i.qty * i.unit_price), 0) AS total_yen
FROM (
  SELECT id, ordered_at, note FROM orders
  WHERE customer_id = 42 AND status = 1
  ORDER BY ordered_at DESC, id DESC
  LIMIT 100
) AS p
LEFT JOIN order_items AS i ON i.order_id = p.id
GROUP BY p.id, p.ordered_at, p.note
ORDER BY p.ordered_at DESC, p.id DESC"""

LOCK_SQL = "SELECT id, note FROM orders WHERE id = 1 FOR UPDATE"
LOCK_HOLD_SEC = 1.0

def connect(database=DB):
    return pymysql.connect(
        host=os.environ["TIDB_HOST"],
        port=int(os.environ.get("TIDB_PORT", "4000")),
        user=os.environ["TIDB_USER"],
        password=os.environ["TIDB_PASSWORD"],
        database=database,
        ssl=ssl.create_default_context(),
        autocommit=True,
    )

def measure(cur, sql):
    """SQL を実行して全行を受け取るまでの秒数と、その SQL の RU を返す。"""
    start = time.perf_counter()
    cur.execute(sql)
    cur.fetchall()
    sec = time.perf_counter() - start
    cur.execute("SELECT @@tidb_last_query_info")
    return sec, json.loads(cur.fetchone()[0])["ru_consumption"]

def setup():
    cur = connect(None).cursor()
    for sql in SETUP:
        cur.execute(sql)
    cur.execute("SELECT COUNT(*) FROM orders")
    if cur.fetchone()[0] == 0:
        for sql in LOAD:
            cur.execute(sql)
    cur.execute("SELECT (SELECT COUNT(*) FROM orders), (SELECT COUNT(*) FROM order_items)")
    print("orders, order_items:", cur.fetchone())

def stats():
    cur = connect().cursor()
    for name, sql in (("A", QUERY_A), ("B", QUERY_B)):
        cur.execute("EXPLAIN " + sql)
        print(name, "stats:pseudo", "stats:pseudo" in str(cur.fetchall()))

def plan(cur, sql):
    """推定行数を除いた実行計画を返す。"""
    cur.execute("EXPLAIN " + sql)
    return [row[:1] + row[2:] for row in cur.fetchall()]

def pair_ab():
    cur = connect().cursor()
    queries = {"A": QUERY_A, "B": QUERY_B}
    return list(queries), lambda cond: measure(cur, queries[cond])

def pair_conc():
    cur = connect().cursor()

    def run(concurrency):
        cur.execute(f"SET SESSION tidb_executor_concurrency = {concurrency}")
        return measure(cur, QUERY_A)

    plans = []
    for concurrency in (1, 5):
        cur.execute(f"SET SESSION tidb_executor_concurrency = {concurrency}")
        plans.append(plan(cur, QUERY_A))
    print("same plan:", plans[0] == plans[1])
    return [1, 5], run

def pair_lock():
    holder = connect().cursor()
    waiter = connect().cursor()

    def run(cond):
        if cond == "wait":
            # 別のセッションが同じ行をロックし、LOCK_HOLD_SEC 秒後に解放する
            holder.execute("BEGIN")
            holder.execute(LOCK_SQL)
            holder.fetchall()
            timer = threading.Timer(LOCK_HOLD_SEC, lambda: holder.execute("ROLLBACK"))
            timer.start()
        waiter.execute("BEGIN")
        result = measure(waiter, LOCK_SQL)
        waiter.execute("ROLLBACK")
        if cond == "wait":
            timer.join()
        return result

    return ["no_wait", "wait"], run

def bench(pair, warmup=3, repeat=20):
    conds, run = {"ab": pair_ab, "conc": pair_conc, "lock": pair_lock}[pair]()
    for _ in range(warmup):
        for cond in conds:
            run(cond)
    results = {cond: [] for cond in conds}
    for i in range(repeat):
        for cond in conds:
            sec, ru = run(cond)
            results[cond].append((sec, ru))
            print(json.dumps({"i": i, "cond": cond, "sec": round(sec, 4), "ru": round(ru, 3)}))
    for cond, rows in results.items():
        for label, values in (("sec", [r[0] for r in rows]), ("ru", [r[1] for r in rows])):
            q1, median, q3 = statistics.quantiles(values, n=4, method="inclusive")
            print(f"{cond} {label}: median {median:.4g} (q1 {q1:.4g}, q3 {q3:.4g})")

if __name__ == "__main__":
    {"setup": setup, "stats": stats, "bench": lambda: bench(sys.argv[2])}[sys.argv[1]]()