用 Python cProfile 排查 CSV 脚本慢：定位热点并验证集合查询优化

python3 --version
python3 -m cProfile --help
mkdir cprofile-csv-demo
cd cprofile-csv-demo

python3 - <<'PY_DATA'
import csv
from pathlib import Path

with open("samples.csv", "x", encoding="utf-8", newline="") as file:
    writer = csv.writer(file)
    writer.writerow(["sample_id", "note"])
    for i in range(20000):
        note = f"record {i}"
        if i % 4999 == 0:
            note = f"记录 {i}, item\nsecond line"
        writer.writerow([f"S{i % 10000:05d}", note])

with Path("allowed.txt").open("x", encoding="utf-8") as file:
    for i in range(0, 10000, 2):
        file.write(f"S{i:05d}\n")
    file.write("S00000\n")
print("Created 20000 records and 5001 allow-list entries")
PY_DATA

"""Profile list versus set membership while filtering a demo CSV."""
import argparse
import csv
from pathlib import Path
import sys


def is_allowed(sample_id, allowed_ids):
    return sample_id in allowed_ids


def filter_samples(source, allow_file, output, index):
    allowed_ids = allow_file.read_text(encoding="utf-8").splitlines()
    if index == "set":
        allowed_ids = set(allowed_ids)
    kept = 0
    with source.open(encoding="utf-8-sig", newline="") as incoming:
        reader = csv.reader(incoming, strict=True)
        if next(reader, None) != ["sample_id", "note"]:
            raise ValueError("expected header: sample_id,note")
        with output.open("x", encoding="utf-8", newline="") as outgoing:
            writer = csv.writer(outgoing)
            writer.writerow(["sample_id", "note"])
            for number, row in enumerate(reader, 1):
                if len(row) != 2 or not row[0]:
                    raise ValueError(f"invalid data record {number}")
                if is_allowed(row[0], allowed_ids):
                    writer.writerow(row)
                    kept += 1
    return kept


def main():
    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
    parser.add_argument("source", type=Path)
    parser.add_argument("allow_file", type=Path)
    parser.add_argument("output", type=Path, help="new output CSV; must not exist")
    parser.add_argument("--index", choices=["list", "set"], default="list")
    args = parser.parse_args()
    try:
        kept = filter_samples(args.source, args.allow_file, args.output, args.index)
    except (OSError, ValueError, csv.Error) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    print(f"kept={kept}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

python3 filter_samples.py --help
python3 -m cProfile -o list.prof filter_samples.py \
  samples.csv allowed.txt list.csv --index list

python3 - <<'PY_STATS'
from pstats import Stats, SortKey

stats = Stats("list.prof").strip_dirs()
print("=== cumulative time ===")
stats.sort_stats(SortKey.CUMULATIVE).print_stats(12)
print("=== function's own time ===")
stats.sort_stats(SortKey.TIME).print_stats(12)
print("=== membership call sites ===")
stats.print_callers("is_allowed")
PY_STATS

python3 -m cProfile -o set.prof filter_samples.py \
  samples.csv allowed.txt set.csv --index set
python3 - <<'PY_COMPARE'
from pstats import Stats, SortKey

for name in ("list.prof", "set.prof"):
    print(f"=== {name} ===")
    Stats(name).strip_dirs().sort_stats(SortKey.TIME).print_stats("is_allowed")
PY_COMPARE

python3 - <<'PY_CHECK'
import csv
from pathlib import Path

assert Path("list.csv").read_bytes() == Path("set.csv").read_bytes()
with open("set.csv", encoding="utf-8", newline="") as file:
    reader = csv.reader(file)
    assert next(reader) == ["sample_id", "note"]
    rows = list(reader)
assert len(rows) == 10000
assert [row[0] for row in rows] == [
    f"S{i % 10000:05d}" for i in range(20000) if i % 2 == 0
]
assert rows[0][1] == "记录 0, item\nsecond line"
assert sum(row[0] == "S00000" for row in rows) == 2
print("OK: identical bytes, 10000 rows, order and repeated IDs preserved")
PY_CHECK

python3 - <<'PY_TIME'
from pathlib import Path
from statistics import median
import subprocess
import sys
import tempfile
from time import perf_counter

with tempfile.TemporaryDirectory(prefix="cprofile-timing-") as directory:
    for index in ("list", "set"):
        elapsed = []
        for run in range(3):
            output = Path(directory) / f"{index}-{run}.csv"
            start = perf_counter()
            subprocess.run([
                sys.executable, "filter_samples.py", "samples.csv",
                "allowed.txt", str(output), "--index", index,
            ], check=True, stdout=subprocess.DEVNULL)
            elapsed.append(perf_counter() - start)
            assert output.read_bytes() == Path("list.csv").read_bytes()
        print(f"{index}: median_seconds={median(elapsed):.6f}")
PY_TIME