Phân tích CSV vượt ngân sách bộ nhớ với DuckDB: đo thời gian và dữ liệu ghi tạm

1python3 -m venv .venv
2. .venv/bin/activate
3python3 -m pip install duckdb==1.4.4
4python3 experiment.py

  1import hashlib
  2import json
  3import platform
  4import resource
  5import subprocess
  6import sys
  7import tempfile
  8from pathlib import Path
  9from time import perf_counter
 10
 11import duckdb
 12
 13
 14ROWS = 900_000
 15LIMITS_MIB = (64, 512, 512, 64)
 16
 17
 18def quote(path):
 19    return "'" + str(path).replace("'", "''") + "'"
 20
 21
 22def make_csv(path):
 23    with path.open("w", encoding="ascii", newline="") as out:
 24        out.write("id,score,payload\n")
 25        for row_id in range(ROWS):
 26            score = (row_id * 104_729 + 12_345) % 100_003
 27            payload = hashlib.shake_128(str(row_id).encode("ascii")).hexdigest(80)
 28            out.write(f"{row_id},{score},{payload}\n")
 29
 30
 31def has_sort(node):
 32    return node.get("operator_type") == "ORDER_BY" or any(
 33        has_sort(child) for child in node.get("children", [])
 34    )
 35
 36
 37def sort_trial(csv, run, limit_mib):
 38    profile_path = run / "profile.json"
 39    output_path = run / "sorted.parquet"
 40    with duckdb.connect() as db:
 41        db.execute("SET threads = 1")
 42        db.execute("SET preserve_insertion_order = true")
 43        db.execute(f"SET memory_limit = '{limit_mib} MiB'")
 44        db.execute(f"SET temp_directory = {quote(run / 'spill')}")
 45        db.execute("SET profiling_coverage = 'ALL'")
 46        db.execute(f"SET profiling_output = {quote(profile_path)}")
 47        db.execute("SET enable_profiling = 'json'")
 48        sql = f"""
 49            COPY (
 50                SELECT id, score, payload
 51                FROM read_csv(
 52                    {quote(csv)}, header=true, auto_detect=false, delim=',',
 53                    columns={{'id': 'BIGINT', 'score': 'INTEGER', 'payload': 'VARCHAR'}},
 54                    max_line_size=2048, buffer_size=32768
 55                )
 56                ORDER BY score, id
 57            ) TO {quote(output_path)} (FORMAT PARQUET, COMPRESSION SNAPPY)
 58        """
 59        started = perf_counter()
 60        written = db.execute(sql).fetchone()[0]
 61        elapsed = perf_counter() - started
 62        rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
 63        rss_bytes = rss if sys.platform == "darwin" else rss * 1024
 64        profile = json.loads(profile_path.read_text(encoding="utf-8"))
 65        assert written == ROWS and has_sort(profile)
 66        print(json.dumps({
 67            "limit_mib": limit_mib,
 68            "elapsed_s": round(elapsed, 4),
 69            "peak_temp_bytes": profile["system_peak_temp_dir_size"],
 70            "peak_buffer_bytes": profile["system_peak_buffer_memory"],
 71            "peak_rss_bytes": rss_bytes,
 72            "parquet_bytes": output_path.stat().st_size,
 73        }))
 74
 75
 76def verify_sorted(path):
 77    digest = hashlib.sha256()
 78    count = id_sum = 0
 79    previous = (-1, -1)
 80    with duckdb.connect() as db:
 81        db.execute("SET threads = 1")
 82        db.execute("SET preserve_insertion_order = true")
 83        cursor = db.execute("SELECT id, score, payload FROM read_parquet(?)", [str(path)])
 84        while batch := cursor.fetchmany(8192):
 85            for row_id, score, payload in batch:
 86                key = (score, row_id)
 87                assert previous < key
 88                assert 0 <= row_id < ROWS and score == (row_id * 104_729 + 12_345) % 100_003
 89                previous = key
 90                count += 1
 91                id_sum += row_id
 92                digest.update(f"{row_id},{score},{payload}\n".encode("ascii"))
 93    assert count == ROWS and id_sum == ROWS * (ROWS - 1) // 2
 94    return count, id_sum, digest.hexdigest()
 95
 96
 97def main():
 98    if sys.platform not in {"darwin", "linux"}:
 99        raise SystemExit("This experiment supports macOS and Linux")
100    if duckdb.__version__ != "1.4.4":
101        raise SystemExit("Install duckdb==1.4.4")
102    print(f"DuckDB {duckdb.__version__}, Python {platform.python_version()}, {platform.system()} {platform.machine()}")
103    with tempfile.TemporaryDirectory(prefix="duckdb-sort-") as folder:
104        workspace = Path(folder)
105        csv = workspace / "source.csv"
106        make_csv(csv)
107        print(f"rows={ROWS}, csv_bytes={csv.stat().st_size}, threads=1")
108        print("run order (MiB):", LIMITS_MIB)
109        runs = []
110        for index, limit_mib in enumerate(LIMITS_MIB, start=1):
111            run = workspace / f"run-{index}"
112            run.mkdir()
113            result = subprocess.check_output(
114                [sys.executable, __file__, "--trial", str(csv), str(run), str(limit_mib)],
115                text=True,
116            )
117            print(f"trial={index} {result.strip()}")
118            runs.append(run / "sorted.parquet")
119        signatures = [verify_sorted(path) for path in runs]
120        assert all(signature == signatures[0] for signature in signatures)
121        print("all sorted outputs (count, sum(id), sha256):", signatures[0])
122
123
124if __name__ == "__main__":
125    if len(sys.argv) == 5 and sys.argv[1] == "--trial":
126        sort_trial(Path(sys.argv[2]), Path(sys.argv[3]), int(sys.argv[4]))
127    else:
128        main()