scripts/benchmark/analyze_histogram_bench.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""
Local benchmark for ANALYZE TABLE histogram algorithms.
The script creates a synthetic FUSE table, runs ANALYZE TABLE WITH HISTOGRAM
with the window, KLL fast, and KLL full algorithms, samples databend-query RSS
while ANALYZE is running, and compares histogram-based cardinality estimates
with exact counts.
It is intended for local experiments and is not part of the CI test suite.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import re
import shlex
import subprocess
import sys
import threading
import time
from collections.abc import Callable
from dataclasses import asdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from scripts.bench_common.bendsql import BendSQL, QueryResult, csv_records
DEFAULT_LOCAL_DSN = "databend://root:@127.0.0.1:38000?sslmode=disable"
HISTOGRAM_RE = re.compile(
r'\[bucket id: (?P<id>\d+), min: "(?P<min>.*?)", max: "(?P<max>.*?)", '
r"ndv: (?P<ndv>[-+0-9.eE]+), count: (?P<count>[-+0-9.eE]+)\]"
)
@dataclass(frozen=True)
class Bucket:
bucket_id: int
lower: float
upper: float
ndv: float
count: float
@dataclass(frozen=True)
class Probe:
column: str
kind: str
lower: int
upper: int | None = None
@property
def label(self) -> str:
if self.kind == "eq":
return f"{self.column} = {self.lower}"
return f"{self.column} BETWEEN {self.lower} AND {self.upper}"
@property
def sql_predicate(self) -> str:
column = quote_ident(self.column)
if self.kind == "eq":
return f"{column} = {self.lower}"
assert self.upper is not None
return f"{column} BETWEEN {self.lower} AND {self.upper}"
@dataclass
class AnalyzeResult:
algorithm: str
elapsed_sec: float
query_pids: list[int]
rss_baseline_kb: int | None
rss_peak_kb: int | None
rss_delta_kb: int | None
@dataclass
class AccuracyResult:
algorithm: str
column: str
distribution: str
predicate_kind: str
predicate: str
exact_count: float
estimated_count: float
absolute_error: float
q_error: float
relative_error: float | None
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Compare ANALYZE histogram algorithm cost and accuracy."
)
parser.add_argument("--dsn", default=os.getenv("BENDSQL_DSN"), help="bendsql DSN")
parser.add_argument("--bendsql", default=os.getenv("BENDSQL", "bendsql"))
parser.add_argument("--database", default="tmp_analyze_histogram_bench")
parser.add_argument("--table", default="t")
parser.add_argument("--rows", type=int, default=5_000_000)
parser.add_argument("--batch-rows", type=int, default=500_000)
parser.add_argument("--domain", type=int, default=100_000)
parser.add_argument(
"--skew-percent",
type=int,
default=80,
help="Percent of rows forced to skew_key = 0.",
)
parser.add_argument("--max-threads", type=int, default=8)
parser.add_argument("--kll-error-rate", type=float, default=0.01)
parser.add_argument(
"--algorithms",
default="window,kll_fast,kll_full",
help="Comma-separated analyze algorithms to run: window,kll_fast,kll_full.",
)
parser.add_argument("--repeat", type=int, default=1)
parser.add_argument("--sample-interval-sec", type=float, default=0.05)
parser.add_argument(
"--run-mode",
choices=("sequential", "isolated-query"),
default="sequential",
help=(
"sequential reuses the current query process; isolated-query starts a fresh "
"databend-query process for setup and for each measured ANALYZE run."
),
)
parser.add_argument(
"--query-bin",
default="target/debug/databend-query",
help="databend-query binary used by --run-mode=isolated-query.",
)
parser.add_argument(
"--query-config",
default="_data/local/databend-query.toml",
help="databend-query config used by --run-mode=isolated-query.",
)
parser.add_argument(
"--query-ready-timeout-sec",
type=float,
default=60.0,
help="Timeout while waiting for a managed databend-query process to accept SQL.",
)
parser.add_argument(
"--query-shutdown-timeout-sec",
type=float,
default=10.0,
help="Grace period before killing a managed databend-query process.",
)
parser.add_argument(
"--query-log-dir",
default="target/bench-results",
help="Directory for managed databend-query stdout/stderr logs.",
)
parser.add_argument(
"--query-pid",
action="append",
type=int,
default=[],
help="databend-query PID to sample. Can be passed multiple times.",
)
parser.add_argument(
"--query-process-substring",
default="databend-query",
help="Used to auto-discover query PIDs when --query-pid is omitted.",
)
parser.add_argument("--reuse-table", action="store_true")
parser.add_argument("--keep", action="store_true", help="Do not drop database after run.")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--output", default=None, help="Write JSON result to this path.")
return parser.parse_args()
def quote_ident(name: str) -> str:
return "`" + name.replace("`", "``") + "`"
def fq_table(database: str, table: str) -> str:
return f"{quote_ident(database)}.{quote_ident(table)}"
def build_client(args: argparse.Namespace) -> BendSQL:
if args.run_mode == "isolated-query" and not args.dsn:
args.dsn = DEFAULT_LOCAL_DSN
return BendSQL(args.bendsql, args.dsn, args.dry_run)
def run_sql(client: BendSQL, sql: str, label: str, echo_stdout: bool = True) -> QueryResult:
return client.run(sql, label, output="csv", echo_stdout=echo_stdout)
def session_sql(args: argparse.Namespace, sql: str) -> str:
return "\n".join(
[
"SET enable_table_snapshot_stats = 1;",
"SET enable_analyze_histogram = 1;",
f"SET max_threads = {args.max_threads};",
sql,
]
)
def prepare_table(client: BendSQL, args: argparse.Namespace) -> None:
table = fq_table(args.database, args.table)
if not args.reuse_table:
run_sql(client, f"DROP DATABASE IF EXISTS {quote_ident(args.database)}", "drop database")
run_sql(client, f"CREATE DATABASE {quote_ident(args.database)}", "create database")
run_sql(
client,
f"""
CREATE TABLE {table} (
id UInt64,
uniform_key UInt64,
skew_key UInt64,
payload String
)
""",
"create table",
)
inserted = 0
while inserted < args.rows:
batch_rows = min(args.batch_rows, args.rows - inserted)
insert_sql = f"""
INSERT INTO {table}
SELECT
number + {inserted} AS id,
(number + {inserted}) % {args.domain} AS uniform_key,
IF(((number + {inserted}) % 100) < {args.skew_percent},
0,
(number + {inserted}) % {args.domain}) AS skew_key,
concat('payload_', to_string((number + {inserted}) % {args.domain})) AS payload
FROM numbers({batch_rows})
"""
run_sql(client, session_sql(args, insert_sql), f"insert batch offset={inserted}", False)
inserted += batch_rows
else:
run_sql(client, f"CREATE DATABASE IF NOT EXISTS {quote_ident(args.database)}", "ensure database")
run_sql(
client,
f"SELECT count(), count(DISTINCT uniform_key), count(DISTINCT skew_key) FROM {table}",
"table cardinality",
)
def discover_query_pids(process_substring: str) -> list[int]:
pids: list[int] = []
self_pid = os.getpid()
proc_dir = Path("/proc")
if not proc_dir.exists():
return pids
for entry in proc_dir.iterdir():
if not entry.name.isdigit():
continue
pid = int(entry.name)
if pid == self_pid:
continue
try:
cmdline = (entry / "cmdline").read_bytes().replace(b"\x00", b" ").decode()
except (OSError, UnicodeDecodeError):
continue
if process_substring in cmdline:
pids.append(pid)
return sorted(pids)
def rss_kb(pid: int) -> int | None:
try:
for line in Path(f"/proc/{pid}/status").read_text().splitlines():
if line.startswith("VmRSS:"):
return int(line.split()[1])
except OSError:
return None
return None
def total_rss_kb(pids: list[int]) -> int | None:
values = [value for pid in pids if (value := rss_kb(pid)) is not None]
if not values:
return None
return sum(values)
def safe_label(label: str) -> str:
return re.sub(r"[^A-Za-z0-9_.-]+", "_", label).strip("_")
def tail_file(path: Path, max_bytes: int = 8192) -> str:
try:
with path.open("rb") as file:
file.seek(0, os.SEEK_END)
size = file.tell()
file.seek(max(0, size - max_bytes))
return file.read().decode(errors="replace")
except OSError:
return ""
def start_query_process(args: argparse.Namespace, label: str) -> tuple[subprocess.Popen[Any], Path]:
if args.dry_run:
raise RuntimeError("managed query process is not available in dry-run mode")
log_dir = Path(args.query_log_dir)
log_dir.mkdir(parents=True, exist_ok=True)
log_path = log_dir / f"analyze_histogram_{safe_label(label)}_{int(time.time() * 1000)}.log"
cmd = [args.query_bin, "-c", args.query_config]
printable = " ".join(shlex.quote(part) for part in cmd)
print(f"\n[start query {label}] {printable} > {log_path}", flush=True)
env = os.environ.copy()
env.setdefault("RUST_BACKTRACE", "1")
with log_path.open("ab") as log_file:
process = subprocess.Popen(cmd, stdout=log_file, stderr=subprocess.STDOUT, env=env)
return process, log_path
def wait_for_query(
client: BendSQL,
args: argparse.Namespace,
process: subprocess.Popen[Any],
log_path: Path,
label: str,
) -> None:
deadline = time.monotonic() + args.query_ready_timeout_sec
cmd = build_bendsql_cmd(client, "SELECT 1")
last_error = ""
while time.monotonic() < deadline:
exit_code = process.poll()
if exit_code is not None:
log_tail = tail_file(log_path)
raise RuntimeError(
f"databend-query exited while starting {label}, exit={exit_code}\n{log_tail}"
)
try:
probe = subprocess.run(cmd, text=True, capture_output=True, timeout=5)
except subprocess.SubprocessError as err:
last_error = str(err)
else:
if probe.returncode == 0:
print(f"[query ready {label}] pid={process.pid}", flush=True)
return
last_error = (probe.stderr or probe.stdout).strip()
time.sleep(0.5)
log_tail = tail_file(log_path)
raise RuntimeError(
f"timed out waiting for databend-query {label}: {last_error}\n{log_tail}"
)
def stop_query_process(
process: subprocess.Popen[Any],
args: argparse.Namespace,
label: str,
) -> None:
if process.poll() is not None:
return
print(f"[stop query {label}] pid={process.pid}", flush=True)
process.terminate()
try:
process.wait(timeout=args.query_shutdown_timeout_sec)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
def with_managed_query(
client: BendSQL,
args: argparse.Namespace,
label: str,
callback: Callable[[list[int]], Any],
) -> Any:
if args.dry_run:
return callback([])
process, log_path = start_query_process(args, label)
try:
wait_for_query(client, args, process, log_path, label)
return callback([process.pid])
finally:
stop_query_process(process, args, label)
class RssMonitor:
def __init__(self, pids: list[int], interval_sec: float) -> None:
self.pids = pids
self.interval_sec = interval_sec
self.baseline_kb: int | None = total_rss_kb(pids)
self.peak_kb: int | None = self.baseline_kb
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, daemon=True)
def __enter__(self) -> "RssMonitor":
self._thread.start()
return self
def __exit__(self, *_: Any) -> None:
self._stop.set()
self._thread.join()
def _run(self) -> None:
while not self._stop.is_set():
current = total_rss_kb(self.pids)
if current is not None:
if self.baseline_kb is None:
self.baseline_kb = current
if self.peak_kb is None or current > self.peak_kb:
self.peak_kb = current
self._stop.wait(self.interval_sec)
@property
def delta_kb(self) -> int | None:
if self.baseline_kb is None or self.peak_kb is None:
return None
return max(0, self.peak_kb - self.baseline_kb)
def build_bendsql_cmd(client: BendSQL, sql: str, output: str = "csv") -> list[str]:
cmd = [client.binary, "--output", output, f"--query={sql}"]
if client.dsn:
cmd[1:1] = ["--dsn", client.dsn]
return cmd
def run_measured_sql(
client: BendSQL,
sql: str,
label: str,
pids: list[int],
sample_interval_sec: float,
) -> tuple[QueryResult, int | None, int | None, int | None]:
cmd = build_bendsql_cmd(client, sql)
printable = " ".join(shlex.quote(part) for part in cmd)
print(f"\n[{label}] {printable}", flush=True)
if client.dry_run:
return QueryResult(label, sql, 0.0, "", ""), None, None, None
start = time.monotonic()
with RssMonitor(pids, sample_interval_sec) as monitor:
proc = subprocess.run(cmd, text=True, capture_output=True)
elapsed = time.monotonic() - start
if proc.returncode != 0:
print(proc.stdout, end="", file=sys.stdout)
print(proc.stderr, end="", file=sys.stderr)
raise RuntimeError(f"bendsql failed for {label}, exit={proc.returncode}")
if proc.stdout.strip():
print(proc.stdout.strip(), flush=True)
if proc.stderr.strip():
print(proc.stderr.strip(), file=sys.stderr, flush=True)
return (
QueryResult(label, sql, elapsed, proc.stdout, proc.stderr),
monitor.baseline_kb,
monitor.peak_kb,
monitor.delta_kb,
)
def analyze_sql(args: argparse.Namespace, algorithm: str) -> str:
options = f"ALGORITHM = '{algorithm}'"
if algorithm in {"kll_fast", "kll_full"}:
options += f", ERROR_RATE = {args.kll_error_rate}"
return session_sql(args, f"ANALYZE TABLE {fq_table(args.database, args.table)} WITH HISTOGRAM {options}")
def run_analyze(
client: BendSQL,
args: argparse.Namespace,
algorithm: str,
iteration: int,
pids: list[int],
) -> AnalyzeResult:
sql = analyze_sql(args, algorithm)
result, baseline_kb, peak_kb, delta_kb = run_measured_sql(
client,
sql,
f"analyze {algorithm} iter={iteration}",
pids,
args.sample_interval_sec,
)
return AnalyzeResult(
algorithm=algorithm,
elapsed_sec=result.elapsed_sec,
query_pids=pids,
rss_baseline_kb=baseline_kb,
rss_peak_kb=peak_kb,
rss_delta_kb=delta_kb,
)
def parse_histogram(histogram: str) -> list[Bucket]:
buckets: list[Bucket] = []
for match in HISTOGRAM_RE.finditer(histogram):
try:
buckets.append(
Bucket(
bucket_id=int(match.group("id")),
lower=float(match.group("min")),
upper=float(match.group("max")),
ndv=float(match.group("ndv")),
count=float(match.group("count")),
)
)
except ValueError:
continue
return buckets
def load_histograms(client: BendSQL, args: argparse.Namespace) -> dict[str, dict[str, Any]]:
result = run_sql(
client,
f"""
SELECT column_name, distinct_count, histogram
FROM fuse_statistic('{args.database}', '{args.table}')
WHERE column_name IN ('uniform_key', 'skew_key')
ORDER BY column_name
""",
"load histograms",
False,
)
histograms: dict[str, dict[str, Any]] = {}
for row in csv_records(result.stdout):
if len(row) < 3:
continue
column_name, distinct_count, histogram = row[0], row[1], row[2]
histograms[column_name] = {
"distinct_count": int(distinct_count) if distinct_count else None,
"histogram": histogram,
"buckets": parse_histogram(histogram),
}
return histograms
def build_probes(domain: int) -> list[Probe]:
domain = max(domain, 10)
small_width = max(1, domain // 100)
medium_width = max(1, domain // 20)
return [
Probe("uniform_key", "eq", 0),
Probe("uniform_key", "eq", domain // 4),
Probe("uniform_key", "eq", domain // 2),
Probe("uniform_key", "eq", domain - 1),
Probe("uniform_key", "range", 0, small_width),
Probe("uniform_key", "range", domain // 3, domain // 3 + small_width),
Probe("uniform_key", "range", domain // 2, domain // 2 + medium_width),
Probe("skew_key", "eq", 0),
Probe("skew_key", "eq", domain // 4),
Probe("skew_key", "range", 0, 0),
Probe("skew_key", "range", 0, small_width),
Probe("skew_key", "range", domain // 2, domain // 2 + medium_width),
]
def probe_distribution(probe: Probe) -> str:
if probe.column == "uniform_key":
return "uniform"
if probe.column == "skew_key":
return "skew"
return probe.column
def exact_counts(client: BendSQL, args: argparse.Namespace, probes: list[Probe]) -> dict[str, float]:
counts: dict[str, float] = {}
table = fq_table(args.database, args.table)
for probe in probes:
result = run_sql(
client,
f"SELECT count() FROM {table} WHERE {probe.sql_predicate}",
f"exact {probe.label}",
False,
)
rows = csv_records(result.stdout)
counts[probe.label] = float(rows[0][0]) if rows and rows[0] else 0.0
return counts
def estimate_probe(buckets: list[Bucket], probe: Probe) -> float:
if probe.kind == "eq":
estimate = 0.0
value = float(probe.lower)
for bucket in buckets:
if bucket.lower <= value <= bucket.upper:
estimate += bucket.count / max(bucket.ndv, 1.0)
return estimate
assert probe.upper is not None
lower = float(probe.lower)
upper = float(probe.upper)
estimate = 0.0
for bucket in buckets:
overlap_lower = max(lower, bucket.lower)
overlap_upper = min(upper, bucket.upper)
if overlap_lower > overlap_upper:
continue
if bucket.lower == bucket.upper:
estimate += bucket.count
continue
bucket_width = bucket.upper - bucket.lower + 1.0
overlap_width = overlap_upper - overlap_lower + 1.0
estimate += bucket.count * max(0.0, min(1.0, overlap_width / bucket_width))
return estimate
def evaluate_accuracy(
algorithm: str,
histograms: dict[str, dict[str, Any]],
probes: list[Probe],
counts: dict[str, float],
) -> list[AccuracyResult]:
results: list[AccuracyResult] = []
for probe in probes:
buckets = histograms.get(probe.column, {}).get("buckets", [])
estimate = estimate_probe(buckets, probe)
exact = counts[probe.label]
absolute = abs(estimate - exact)
q_error = cardinality_q_error(estimate, exact)
relative = None if exact == 0 else absolute / exact
results.append(
AccuracyResult(
algorithm=algorithm,
column=probe.column,
distribution=probe_distribution(probe),
predicate_kind=probe.kind,
predicate=probe.label,
exact_count=exact,
estimated_count=estimate,
absolute_error=absolute,
q_error=q_error,
relative_error=relative,
)
)
return results
def cardinality_q_error(estimated: float, exact: float) -> float:
estimated = max(estimated, 1.0)
exact = max(exact, 1.0)
return max(estimated / exact, exact / estimated)
def summarize_histograms(histograms: dict[str, dict[str, Any]]) -> dict[str, Any]:
summary: dict[str, Any] = {}
for column, info in histograms.items():
buckets: list[Bucket] = info["buckets"]
summary[column] = {
"distinct_count": info["distinct_count"],
"num_buckets": len(buckets),
"bucket_count_sum": sum(bucket.count for bucket in buckets),
"bucket_ndv_sum": sum(bucket.ndv for bucket in buckets),
"histogram_bytes": len(info["histogram"]),
}
return summary
def summarize_accuracy_values(values: list[AccuracyResult]) -> dict[str, Any]:
q_errors = sorted(result.q_error for result in values)
relatives = sorted(
result.relative_error for result in values if result.relative_error is not None
)
absolutes = sorted(result.absolute_error for result in values)
return {
"probe_count": len(values),
"q_error_count": len(q_errors),
"max_q_error": max(q_errors) if q_errors else None,
"mean_q_error": sum(q_errors) / len(q_errors) if q_errors else None,
"geomean_q_error": geometric_mean(q_errors),
"p50_q_error": percentile(q_errors, 0.50),
"p90_q_error": percentile(q_errors, 0.90),
"p95_q_error": percentile(q_errors, 0.95),
"relative_probe_count": len(relatives),
"max_absolute_error": max(absolutes) if absolutes else None,
"mean_absolute_error": sum(absolutes) / len(absolutes) if absolutes else None,
"max_relative_error": max(relatives) if relatives else None,
"mean_relative_error": sum(relatives) / len(relatives) if relatives else None,
}
def geometric_mean(values: list[float]) -> float | None:
if not values:
return None
return math.exp(sum(math.log(value) for value in values) / len(values))
def percentile(values: list[float], quantile: float) -> float | None:
if not values:
return None
if len(values) == 1:
return values[0]
index = (len(values) - 1) * quantile
lower_index = int(index)
upper_index = min(lower_index + 1, len(values) - 1)
fraction = index - lower_index
return values[lower_index] * (1.0 - fraction) + values[upper_index] * fraction
def summarize_accuracy(results: list[AccuracyResult]) -> dict[str, Any]:
by_algorithm: dict[str, list[AccuracyResult]] = {}
for result in results:
by_algorithm.setdefault(result.algorithm, []).append(result)
summary: dict[str, Any] = {}
for algorithm, values in by_algorithm.items():
summary[algorithm] = summarize_accuracy_values(values)
return summary
def summarize_accuracy_by_distribution(results: list[AccuracyResult]) -> dict[str, Any]:
by_group: dict[str, dict[str, list[AccuracyResult]]] = {}
for result in results:
by_group.setdefault(result.algorithm, {}).setdefault(result.distribution, []).append(result)
return {
algorithm: {
distribution: summarize_accuracy_values(values)
for distribution, values in sorted(distributions.items())
}
for algorithm, distributions in sorted(by_group.items())
}
def summarize_accuracy_by_kind(results: list[AccuracyResult]) -> dict[str, Any]:
by_group: dict[str, dict[str, list[AccuracyResult]]] = {}
for result in results:
by_group.setdefault(result.algorithm, {}).setdefault(result.predicate_kind, []).append(
result
)
return {
algorithm: {
predicate_kind: summarize_accuracy_values(values)
for predicate_kind, values in sorted(kinds.items())
}
for algorithm, kinds in sorted(by_group.items())
}
def summarize_accuracy_by_distribution_and_kind(
results: list[AccuracyResult],
) -> dict[str, Any]:
by_group: dict[str, dict[str, dict[str, list[AccuracyResult]]]] = {}
for result in results:
by_group.setdefault(result.algorithm, {}).setdefault(result.distribution, {}).setdefault(
result.predicate_kind, []
).append(result)
return {
algorithm: {
distribution: {
predicate_kind: summarize_accuracy_values(values)
for predicate_kind, values in sorted(kinds.items())
}
for distribution, kinds in sorted(distributions.items())
}
for algorithm, distributions in sorted(by_group.items())
}
def main() -> None:
args = parse_args()
if args.rows <= 0 or args.batch_rows <= 0 or args.domain <= 0:
raise ValueError("--rows, --batch-rows, and --domain must be greater than zero")
if not 0 <= args.skew_percent <= 100:
raise ValueError("--skew-percent must be between 0 and 100")
algorithms = [item.strip() for item in args.algorithms.split(",") if item.strip()]
for algorithm in algorithms:
if algorithm not in {"window", "kll_fast", "kll_full"}:
raise ValueError(f"unsupported algorithm: {algorithm}")
probes = build_probes(args.domain)
client = build_client(args)
analyze_results: list[AnalyzeResult] = []
accuracy_results: list[AccuracyResult] = []
histogram_summaries: dict[str, Any] = {}
sampled_pids: list[int] = []
if args.run_mode == "sequential":
sampled_pids = args.query_pid or discover_query_pids(args.query_process_substring)
if sampled_pids:
print(f"Sampling databend-query RSS from pids: {sampled_pids}", flush=True)
else:
print(
"No databend-query PID discovered; RSS fields will be null. "
"Pass --query-pid to enable memory sampling.",
file=sys.stderr,
flush=True,
)
prepare_table(client, args)
counts = exact_counts(client, args, probes)
for iteration in range(1, args.repeat + 1):
for algorithm in algorithms:
analyze_result = run_analyze(client, args, algorithm, iteration, sampled_pids)
analyze_results.append(analyze_result)
histograms = load_histograms(client, args)
key = f"{algorithm}_iter_{iteration}"
histogram_summaries[key] = summarize_histograms(histograms)
accuracy_results.extend(evaluate_accuracy(algorithm, histograms, probes, counts))
else:
if args.query_pid:
print(
"--query-pid is ignored in --run-mode=isolated-query; managed query "
"processes are sampled by their own PIDs.",
file=sys.stderr,
flush=True,
)
def setup(_: list[int]) -> dict[str, float]:
prepare_table(client, args)
return exact_counts(client, args, probes)
counts = with_managed_query(client, args, "setup", setup)
for iteration in range(1, args.repeat + 1):
for algorithm in algorithms:
def run_once(pids: list[int]) -> None:
analyze_result = run_analyze(client, args, algorithm, iteration, pids)
analyze_results.append(analyze_result)
histograms = load_histograms(client, args)
key = f"{algorithm}_iter_{iteration}"
histogram_summaries[key] = summarize_histograms(histograms)
accuracy_results.extend(
evaluate_accuracy(algorithm, histograms, probes, counts)
)
with_managed_query(client, args, f"{algorithm}_iter_{iteration}", run_once)
result = {
"config": {
"database": args.database,
"table": args.table,
"rows": args.rows,
"batch_rows": args.batch_rows,
"domain": args.domain,
"skew_percent": args.skew_percent,
"max_threads": args.max_threads,
"kll_error_rate": args.kll_error_rate,
"algorithms": algorithms,
"repeat": args.repeat,
"run_mode": args.run_mode,
"query_bin": args.query_bin if args.run_mode == "isolated-query" else None,
"query_config": args.query_config if args.run_mode == "isolated-query" else None,
"query_pids": sampled_pids,
},
"analyze_results": [asdict(result) for result in analyze_results],
"histogram_summaries": histogram_summaries,
"accuracy_results": [asdict(result) for result in accuracy_results],
"accuracy_summary": summarize_accuracy(accuracy_results),
"accuracy_summary_by_distribution": summarize_accuracy_by_distribution(accuracy_results),
"accuracy_summary_by_kind": summarize_accuracy_by_kind(accuracy_results),
"accuracy_summary_by_distribution_and_kind": summarize_accuracy_by_distribution_and_kind(
accuracy_results
),
}
print(json.dumps(result, indent=2, sort_keys=True), flush=True)
if args.output:
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
Path(args.output).write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
print(f"Wrote {args.output}", flush=True)
if not args.keep:
if args.run_mode == "sequential":
run_sql(client, f"DROP DATABASE IF EXISTS {quote_ident(args.database)}", "drop database")
else:
def cleanup(_: list[int]) -> None:
run_sql(
client,
f"DROP DATABASE IF EXISTS {quote_ident(args.database)}",
"drop database",
)
with_managed_query(client, args, "cleanup", cleanup)
if __name__ == "__main__":
main()