From 884f86f64ead36b2e6a3ffa10eb415c7475981f5 Mon Sep 17 00:00:00 2001 From: Greg Burd Date: Thu, 9 Jul 2026 12:34:56 -0400 Subject: [PATCH v8 2/3] [NOT FOR MERGE] Add atomics benchmark harness Python-based A/B benchmark harness for comparing the stdatomic.h and traditional atomic implementations under load. Local development tooling; not part of the submitted series. (cherry picked from commit 4d9bcf693a397ccf85ececd4c070bc555bf14efb) (cherry picked from commit ba8192a458f639a5ca7f68e3e9976a4f43dc329b) (cherry picked from commit 3cca1e9c1054e2daedf8d54b938f1c4fadd89342) (cherry picked from commit bfd22b8087b6f6aa3527433988ca4cd20db6c191) (cherry picked from commit 2c1e09739dafb18750f9a2632b0d2db5844e1c5f) (cherry picked from commit bcd153779a98a9c03b1706a62dd0bb773622fc76) --- src/test/benchmarks/README.md | 282 ++++++++++++ src/test/benchmarks/atomics_analyzer.py | 295 ++++++++++++ src/test/benchmarks/atomics_benchmark.py | 553 +++++++++++++++++++++++ src/test/benchmarks/atomics_cli.py | 354 +++++++++++++++ src/test/benchmarks/atomics_config.py | 275 +++++++++++ src/test/benchmarks/atomics_workload.py | 304 +++++++++++++ src/test/benchmarks/data_generator.py | 409 +++++++++++++++++ 7 files changed, 2472 insertions(+) create mode 100644 src/test/benchmarks/README.md create mode 100644 src/test/benchmarks/atomics_analyzer.py create mode 100644 src/test/benchmarks/atomics_benchmark.py create mode 100644 src/test/benchmarks/atomics_cli.py create mode 100644 src/test/benchmarks/atomics_config.py create mode 100644 src/test/benchmarks/atomics_workload.py create mode 100644 src/test/benchmarks/data_generator.py diff --git a/src/test/benchmarks/README.md b/src/test/benchmarks/README.md new file mode 100644 index 00000000000..c4b0c0df195 --- /dev/null +++ b/src/test/benchmarks/README.md @@ -0,0 +1,282 @@ +# PostgreSQL Atomics Performance Benchmark Suite + +This benchmark suite compares the performance of traditional platform-specific atomics implementation versus the C11 stdatomic.h implementation. + +## Overview + +The benchmark evaluates atomic operations impact across different: +- **Workload patterns**: Full scans, aggregations, index scans, concurrent operations +- **Table schemas**: Narrow (4 cols), medium (11 cols), wide (55 cols) +- **Data distributions**: Random, clustered, low cardinality +- **Data scales**: 10K, 100K, 1M, 10M, 100M rows + +## Prerequisites + +1. **Two PostgreSQL builds**: + ```bash + # Traditional atomics + meson setup build-traditional -Duse_stdatomic=no --buildtype=debugoptimized + meson compile -C build-traditional + + # Stdatomic.h implementation + meson setup build-stdatomic -Duse_stdatomic=yes --buildtype=debugoptimized + meson compile -C build-stdatomic + ``` + +2. **Python dependencies**: + ```bash + pip install asyncpg + ``` + +## Usage + +### Quick Start + +Run with default settings (medium test matrix): + +```bash +cd src/test/benchmarks +python -m atomics_cli +``` + +This will: +- Start two PostgreSQL instances (ports 5433 and 5434) +- Run benchmarks on both implementations +- Compare results +- Save to `atomics_benchmark_results/results.json` + +### Full Matrix + +Run comprehensive tests including large datasets: + +```bash +python -m atomics_cli --full-matrix +``` + +**Warning**: This can take several hours. + +### Custom Configuration + +```bash +# Specify custom build directories +python -m atomics_cli \ + --traditional-build /path/to/build-traditional \ + --stdatomic-build /path/to/build-stdatomic + +# Test specific schema and row counts +python -m atomics_cli --schema medium --rows 100000 1000000 + +# Test specific query patterns +python -m atomics_cli --pattern full_scan --pattern aggregation + +# Adjust measurement precision +python -m atomics_cli --warmup 5 --iterations 20 + +# Verbose logging +python -m atomics_cli -v +``` + +### Analyzing Results + +After running benchmarks, analyze the results: + +```bash +python -m atomics_analyzer atomics_benchmark_results/results.json +``` + +Or save to a file: + +```bash +python -m atomics_analyzer atomics_benchmark_results/results.json \ + --output analysis_report.txt +``` + +## Benchmark Methodology + +### Query Patterns + +1. **full_scan**: Sequential scan - stresses spinlock acquisition and buffer management +2. **column_projection**: Selective scan - tuple visibility checks with atomic reads +3. **filtered_scan**: Index + sequential - mixed contention patterns +4. **aggregation**: Heavy computation - shared buffer access +5. **group_by**: Hash aggregation - memory ordering stress +6. **index_scan**: Index-only - buffer pin/unpin contention + +### Measurement Process + +For each scenario: +1. **Warmup**: Run query N times (default: 3) to prime caches +2. **Measure**: Run query M times (default: 10), record each timing +3. **Report**: Use **median** timing to minimize outlier impact + +### Statistical Significance + +- **Noise threshold**: +/-2% (measurements within this are considered equivalent) +- **Significant improvement**: >2% faster +- **Significant regression**: >2% slower + +### Good Benchmarking Practices + +1. **Dedicated hardware**: Run on isolated system, no background processes +2. **CPU frequency**: Disable CPU frequency scaling: + ```bash + echo performance | sudo tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor + ``` +3. **NUMA**: Pin PostgreSQL to single NUMA node if applicable +4. **Disk I/O**: Use fast storage (SSD/NVMe) or tmpfs for data directories +5. **Multiple runs**: Run benchmark multiple times, compare results +6. **Baseline**: Establish baseline performance before code changes + +## Output Format + +Results are saved in JSON format: + +```json +{ + "summary": { + "total_scenarios": 54, + "median_speedup": 1.002, + "mean_speedup": 1.005, + "significant_improvements": 5, + "significant_regressions": 2, + "per_pattern_avg_speedup": { + "full_scan": 0.998, + "aggregation": 1.015, + ... + } + }, + "results": [ + { + "schema": "atomics_medium", + "row_count": 100000, + "distribution": "random", + "pattern": "full_scan", + "traditional_ms": 125.3, + "stdatomic_ms": 126.1, + "speedup": 0.994 + }, + ... + ] +} +``` + +## Interpreting Results + +### Speedup Metric + +- **Speedup = traditional_ms / stdatomic_ms** +- `speedup > 1.0`: stdatomic is **faster** +- `speedup < 1.0`: stdatomic is **slower** +- `speedup ~= 1.0`: Performance is **equivalent** + +### Example Analysis + +``` +VERDICT: NEUTRAL - Performance is statistically equivalent + +Median speedup: 1.002x (stdatomic ~0.2% faster) +Within measurement noise: 47/54 scenarios (87%) +Significant improvements: 5 (9%) +Significant regressions: 2 (4%) +``` + +**Interpretation**: The stdatomic.h implementation has no material performance impact. + +## Troubleshooting + +### Port Already in Use + +``` +ERROR: Port 5433 already in use +``` + +**Solution**: Stop existing PostgreSQL instances or use different ports: +```bash +python -m atomics_cli --traditional-port 5435 --stdatomic-port 5436 +``` + +### Build Directory Not Found + +``` +ERROR: Traditional build directory not found: ./build-traditional +``` + +**Solution**: Build PostgreSQL first or specify correct path: +```bash +python -m atomics_cli --traditional-build /path/to/build +``` + +### Insufficient Disk Space + +For large datasets, ensure adequate space in `/tmp`: +```bash +df -h /tmp +``` + +Consider using custom data directories: +```bash +python -m atomics_cli \ + --traditional-datadir /mnt/fast-storage/pgdata-trad \ + --stdatomic-datadir /mnt/fast-storage/pgdata-std +``` + +### Connection Timeout + +If PostgreSQL takes >30s to start, check: +1. Disk I/O performance +2. System resources (RAM, CPU) +3. PostgreSQL logs in data directory + +## Architecture + +``` +src/test/benchmarks/ +|-- atomics_config.py - Configuration and schema definitions +|-- atomics_benchmark.py - Main orchestrator +|-- atomics_workload.py - Query generation and execution +|-- atomics_cli.py - Command-line interface +|-- atomics_analyzer.py - Results analysis +|-- data_generator.py - Data generation +`-- README_ATOMICS.md - This file +``` + +## For PostgreSQL Committers + +This benchmark suite is designed for: +- Validating atomics implementation changes +- Performance regression testing +- Platform-specific optimization evaluation +- Memory ordering impact analysis + +### CI Integration + +For automated testing: + +```bash +# Quick smoke test (fast, ~5 minutes) +python -m atomics_cli --schema narrow --rows 10000 --iterations 3 + +# Standard test (moderate, ~30 minutes) +python -m atomics_cli --schema medium --rows 100000 1000000 + +# Full validation (slow, several hours) +python -m atomics_cli --full-matrix +``` + +### Expected Performance + +On x86_64 with strong memory model: +- **Median difference**: +/-1-2% (within noise) +- **Pattern variation**: Some patterns may show +/-5% variation +- **Scale impact**: Minimal difference across scales + +On ARM64/RISC-V with weak memory model: +- **Potential regression**: 2-10% (stronger ordering overhead) +- **Workload dependent**: Write-heavy workloads more impacted +- **Trade-off**: Correctness vs. performance + +## References + +- [SEMANTIC_DIFFERENCES.md](../../SEMANTIC_DIFFERENCES.md) - Implementation differences +- [MEMORY_ORDERING_ANALYSIS.md](../../MEMORY_ORDERING_ANALYSIS.md) - Memory ordering decisions +- [PostgreSQL atomics documentation](https://www.postgresql.org/docs/current/atomics.html) diff --git a/src/test/benchmarks/atomics_analyzer.py b/src/test/benchmarks/atomics_analyzer.py new file mode 100644 index 00000000000..3250e894942 --- /dev/null +++ b/src/test/benchmarks/atomics_analyzer.py @@ -0,0 +1,295 @@ +""" +Results analyzer for atomics benchmark comparisons. + +This module provides statistical analysis and visualization of benchmark results, +helping identify performance regressions or improvements from the stdatomic.h +implementation. +""" + +import json +import statistics +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Tuple + + +@dataclass +class ResultComparison: + """Comparison of a single benchmark scenario.""" + schema: str + row_count: int + distribution: str + pattern: str + traditional_ms: float + stdatomic_ms: float + speedup: float + difference_pct: float # (traditional - stdatomic) / traditional * 100 + + +class AtomicsResultAnalyzer: + """Analyzes and compares benchmark results.""" + + def __init__(self, results_file: Path): + self.results_file = results_file + self.comparisons: List[ResultComparison] = [] + self._load_results() + + def _load_results(self): + """Load results from JSON file.""" + with open(self.results_file) as f: + data = json.load(f) + + for r in data.get("results", []): + diff_pct = ( + (r["traditional_ms"] - r["stdatomic_ms"]) / r["traditional_ms"] * 100 + ) + comp = ResultComparison( + schema=r["schema"], + row_count=r["row_count"], + distribution=r["distribution"], + pattern=r["pattern"], + traditional_ms=r["traditional_ms"], + stdatomic_ms=r["stdatomic_ms"], + speedup=r["speedup"], + difference_pct=diff_pct, + ) + self.comparisons.append(comp) + + def analyze_overall(self) -> Dict: + """Analyze overall performance characteristics.""" + if not self.comparisons: + return {} + + speedups = [c.speedup for c in self.comparisons] + diff_pcts = [c.difference_pct for c in self.comparisons] + + # Statistical significance: count how many are outside +/-2% threshold + significant_improvements = sum(1 for d in diff_pcts if d > 2.0) + significant_regressions = sum(1 for d in diff_pcts if d < -2.0) + within_noise = sum(1 for d in diff_pcts if -2.0 <= d <= 2.0) + + return { + "total_scenarios": len(self.comparisons), + "median_speedup": statistics.median(speedups), + "mean_speedup": statistics.mean(speedups), + "stdev_speedup": statistics.stdev(speedups) if len(speedups) > 1 else 0.0, + "min_speedup": min(speedups), + "max_speedup": max(speedups), + "median_diff_pct": statistics.median(diff_pcts), + "mean_diff_pct": statistics.mean(diff_pcts), + "significant_improvements": significant_improvements, + "significant_regressions": significant_regressions, + "within_noise": within_noise, + } + + def analyze_by_pattern(self) -> Dict[str, Dict]: + """Break down performance by query pattern.""" + pattern_groups = {} + for comp in self.comparisons: + if comp.pattern not in pattern_groups: + pattern_groups[comp.pattern] = [] + pattern_groups[comp.pattern].append(comp) + + results = {} + for pattern, comps in pattern_groups.items(): + speedups = [c.speedup for c in comps] + diff_pcts = [c.difference_pct for c in comps] + + results[pattern] = { + "count": len(comps), + "median_speedup": statistics.median(speedups), + "mean_speedup": statistics.mean(speedups), + "median_diff_pct": statistics.median(diff_pcts), + "mean_diff_pct": statistics.mean(diff_pcts), + "improvements": sum(1 for d in diff_pcts if d > 2.0), + "regressions": sum(1 for d in diff_pcts if d < -2.0), + } + + return results + + def analyze_by_scale(self) -> Dict[int, Dict]: + """Break down performance by data scale.""" + scale_groups = {} + for comp in self.comparisons: + if comp.row_count not in scale_groups: + scale_groups[comp.row_count] = [] + scale_groups[comp.row_count].append(comp) + + results = {} + for row_count, comps in scale_groups.items(): + speedups = [c.speedup for c in comps] + diff_pcts = [c.difference_pct for c in comps] + + results[row_count] = { + "count": len(comps), + "median_speedup": statistics.median(speedups), + "mean_speedup": statistics.mean(speedups), + "median_diff_pct": statistics.median(diff_pcts), + "improvements": sum(1 for d in diff_pcts if d > 2.0), + "regressions": sum(1 for d in diff_pcts if d < -2.0), + } + + return results + + def find_regressions(self, threshold_pct: float = 2.0) -> List[ResultComparison]: + """Find scenarios where stdatomic is slower than threshold.""" + return [c for c in self.comparisons if c.difference_pct < -threshold_pct] + + def find_improvements(self, threshold_pct: float = 2.0) -> List[ResultComparison]: + """Find scenarios where stdatomic is faster than threshold.""" + return [c for c in self.comparisons if c.difference_pct > threshold_pct] + + def generate_report(self, output_file: Optional[Path] = None) -> str: + """Generate a comprehensive text report.""" + lines = [] + lines.append("=" * 80) + lines.append("PostgreSQL Atomics Benchmark Analysis Report") + lines.append("=" * 80) + lines.append("") + + # Overall analysis + overall = self.analyze_overall() + lines.append("OVERALL PERFORMANCE") + lines.append("-" * 80) + lines.append(f"Total scenarios: {overall['total_scenarios']}") + lines.append(f"Median speedup: {overall['median_speedup']:.3f}x") + lines.append(f"Mean speedup: {overall['mean_speedup']:.3f}x +/- {overall['stdev_speedup']:.3f}") + lines.append(f"Range: {overall['min_speedup']:.3f}x - {overall['max_speedup']:.3f}x") + lines.append("") + lines.append(f"Median difference: {overall['median_diff_pct']:+.2f}%") + lines.append(f"Mean difference: {overall['mean_diff_pct']:+.2f}%") + lines.append("") + lines.append(f"Significant improvements: {overall['significant_improvements']} (>{2.0:.1f}%)") + lines.append(f"Significant regressions: {overall['significant_regressions']} (<-{2.0:.1f}%)") + lines.append(f"Within measurement noise: {overall['within_noise']} (+/-{2.0:.1f}%)") + lines.append("") + + # Verdict + median_diff = overall["median_diff_pct"] + if abs(median_diff) < 2.0: + verdict = "NEUTRAL - Performance is statistically equivalent" + elif median_diff > 0: + verdict = f"POSITIVE - Stdatomic is faster by {median_diff:.2f}%" + else: + verdict = f"NEGATIVE - Stdatomic is slower by {abs(median_diff):.2f}%" + + lines.append(f"VERDICT: {verdict}") + lines.append("") + + # By pattern + by_pattern = self.analyze_by_pattern() + lines.append("PERFORMANCE BY QUERY PATTERN") + lines.append("-" * 80) + lines.append(f"{'Pattern':<30} {'Median':>10} {'Mean':>10} {'Diff%':>10} {'+/-':>10}") + lines.append("-" * 80) + + for pattern in sorted(by_pattern.keys()): + stats = by_pattern[pattern] + indicator = "+++" if stats["mean_diff_pct"] > 5.0 else " " + indicator = "---" if stats["mean_diff_pct"] < -5.0 else indicator + lines.append( + f"{pattern:<30} {stats['median_speedup']:>9.3f}x " + f"{stats['mean_speedup']:>9.3f}x " + f"{stats['mean_diff_pct']:>9.2f}% " + f"{indicator}" + ) + + lines.append("") + + # By scale + by_scale = self.analyze_by_scale() + lines.append("PERFORMANCE BY DATA SCALE") + lines.append("-" * 80) + lines.append(f"{'Rows':<15} {'Median':>10} {'Mean':>10} {'Diff%':>10} {'Status':<20}") + lines.append("-" * 80) + + for row_count in sorted(by_scale.keys()): + stats = by_scale[row_count] + status = "Faster" if stats["median_diff_pct"] > 2.0 else "Neutral" + status = "Slower" if stats["median_diff_pct"] < -2.0 else status + lines.append( + f"{row_count:<15,} {stats['median_speedup']:>9.3f}x " + f"{stats['mean_speedup']:>9.3f}x " + f"{stats['median_diff_pct']:>9.2f}% " + f"{status:<20}" + ) + + lines.append("") + + # Top regressions + regressions = sorted( + self.find_regressions(2.0), key=lambda c: c.difference_pct + ) + if regressions: + lines.append("TOP REGRESSIONS (stdatomic slower by >2%)") + lines.append("-" * 80) + for comp in regressions[:10]: + lines.append( + f" {comp.difference_pct:>6.2f}% - {comp.pattern:<25} " + f"{comp.schema}/{comp.row_count:,}/{comp.distribution}" + ) + lines.append("") + + # Top improvements + improvements = sorted( + self.find_improvements(2.0), key=lambda c: c.difference_pct, reverse=True + ) + if improvements: + lines.append("TOP IMPROVEMENTS (stdatomic faster by >2%)") + lines.append("-" * 80) + for comp in improvements[:10]: + lines.append( + f" {comp.difference_pct:>6.2f}% - {comp.pattern:<25} " + f"{comp.schema}/{comp.row_count:,}/{comp.distribution}" + ) + lines.append("") + + lines.append("=" * 80) + + report = "\n".join(lines) + + if output_file: + output_file.parent.mkdir(parents=True, exist_ok=True) + with open(output_file, "w") as f: + f.write(report) + + return report + + +def main(): + """CLI entry point for results analysis.""" + import argparse + + parser = argparse.ArgumentParser( + description="Analyze atomics benchmark results" + ) + parser.add_argument( + "results_file", + type=Path, + help="Path to results.json file", + ) + parser.add_argument( + "--output", + "-o", + type=Path, + help="Output file for report (default: stdout)", + ) + + args = parser.parse_args() + + if not args.results_file.exists(): + print(f"ERROR: Results file not found: {args.results_file}") + return 1 + + analyzer = AtomicsResultAnalyzer(args.results_file) + report = analyzer.generate_report(args.output) + + if not args.output: + print(report) + + return 0 + + +if __name__ == "__main__": + import sys + sys.exit(main()) diff --git a/src/test/benchmarks/atomics_benchmark.py b/src/test/benchmarks/atomics_benchmark.py new file mode 100644 index 00000000000..5f5207599ad --- /dev/null +++ b/src/test/benchmarks/atomics_benchmark.py @@ -0,0 +1,553 @@ +""" +Main benchmark orchestrator for comparing traditional vs stdatomic.h implementations. + +This module coordinates: +- PostgreSQL instance lifecycle management +- Schema creation and data loading +- Workload execution on both implementations +- Results collection and comparison +""" + +import asyncio +import asyncpg +import logging +import os +import shutil +import signal +import subprocess +import time +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from .atomics_config import ( + AtomicsBenchmarkConfig, + DataDistribution, + PostgresInstance, + QueryPattern, + TableSchema, +) +from .data_generator import DataGenerator +from .atomics_workload import AtomicsWorkloadRunner, WorkloadResult + +logger = logging.getLogger(__name__) + + +@dataclass +class BenchmarkMetrics: + """Collected metrics for a benchmark run.""" + implementation: str # "traditional" or "stdatomic" + schema_name: str + row_count: int + distribution: str + load_time_seconds: float + table_size_bytes: int + index_size_bytes: int + shared_buffers_hit_ratio: float = 0.0 + cache_hit_ratio: float = 0.0 + + +@dataclass +class ComparisonResult: + """Result of comparing traditional vs stdatomic for one scenario.""" + schema_name: str + row_count: int + distribution: str + pattern: str + traditional_ms: float + stdatomic_ms: float + speedup: float # stdatomic_ms / traditional_ms (< 1.0 means faster) + traditional_metrics: BenchmarkMetrics + stdatomic_metrics: BenchmarkMetrics + + +class PostgresManager: + """Manages PostgreSQL instance lifecycle.""" + + def __init__(self, instance: PostgresInstance): + self.instance = instance + self.process: Optional[subprocess.Popen] = None + self.pool: Optional[asyncpg.Pool] = None + + async def initialize_cluster(self): + """Initialize a new database cluster if needed.""" + if self.instance.data_dir.exists(): + logger.info( + "Data directory %s exists, skipping initdb", self.instance.data_dir + ) + return + + logger.info("Initializing database cluster at %s", self.instance.data_dir) + self.instance.data_dir.mkdir(parents=True, exist_ok=True) + + # For meson build directories (not installed), create symlink for postgres binary + # in initdb directory so initdb can find it + if not self.instance._is_install_dir(): + initdb_dir = self.instance.initdb_bin.parent + postgres_link = initdb_dir / "postgres" + if not postgres_link.exists(): + postgres_link.symlink_to(self.instance.postgres_bin) + logger.debug("Created postgres symlink: %s -> %s", postgres_link, self.instance.postgres_bin) + + cmd = [ + str(self.instance.initdb_bin), + "-D", + str(self.instance.data_dir), + "--no-locale", + "--encoding=UTF8", + ] + + # Set LD_LIBRARY_PATH for installed directories + env = None + if self.instance._is_install_dir(): + lib_dir = self.instance.build_dir / "lib64" + env = os.environ.copy() + env["LD_LIBRARY_PATH"] = f"{lib_dir}:{env.get('LD_LIBRARY_PATH', '')}" + logger.debug("Setting LD_LIBRARY_PATH=%s", env["LD_LIBRARY_PATH"]) + + result = subprocess.run(cmd, capture_output=True, text=True, env=env) + if result.returncode != 0: + raise RuntimeError(f"initdb failed: {result.stderr}") + + logger.info("Database cluster initialized successfully") + + # Configure for benchmarking + self._configure_for_performance() + + def _configure_for_performance(self): + """Update postgresql.conf for benchmark performance.""" + conf_path = self.instance.data_dir / "postgresql.conf" + + settings = { + "shared_buffers": "1GB", + "effective_cache_size": "4GB", + "maintenance_work_mem": "256MB", + "checkpoint_completion_target": "0.9", + "wal_buffers": "16MB", + "default_statistics_target": "100", + "random_page_cost": "1.1", + "effective_io_concurrency": "200", + "work_mem": "32MB", + "min_wal_size": "1GB", + "max_wal_size": "4GB", + "max_worker_processes": "8", + "max_parallel_workers_per_gather": "4", + "max_parallel_workers": "8", + "max_parallel_maintenance_workers": "4", + # Enable timing + "track_activities": "on", + "track_counts": "on", + "track_io_timing": "on", + "track_functions": "all", + # Reduce checkpointing during benchmark + "checkpoint_timeout": "30min", + } + + with open(conf_path, "a") as f: + f.write("\n\n# Performance settings for benchmarking\n") + for key, value in settings.items(): + f.write(f"{key} = {value}\n") + + logger.info("Configured PostgreSQL for performance benchmarking") + + async def start(self): + """Start the PostgreSQL instance.""" + if self.process is not None: + logger.warning("PostgreSQL is already running") + return + + logger.info( + "Starting PostgreSQL (%s) on port %d", + self.instance.name, + self.instance.port, + ) + + cmd = [ + str(self.instance.postgres_bin), + "-D", + str(self.instance.data_dir), + "-p", + str(self.instance.port), + "-c", + "logging_collector=off", + ] + + # Set LD_LIBRARY_PATH for installed directories + env = None + if self.instance._is_install_dir(): + lib_dir = self.instance.build_dir / "lib64" + env = os.environ.copy() + env["LD_LIBRARY_PATH"] = f"{lib_dir}:{env.get('LD_LIBRARY_PATH', '')}" + + self.process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + env=env, + ) + + # Wait for server to be ready + await self._wait_for_ready() + + # Create connection pool + await self._create_pool() + + logger.info("PostgreSQL (%s) started successfully", self.instance.name) + + async def _wait_for_ready(self, timeout: int = 30): + """Wait for PostgreSQL to be ready to accept connections.""" + start = time.time() + while time.time() - start < timeout: + try: + conn = await asyncpg.connect( + host=self.instance.host, + port=self.instance.port, + user=self.instance.user or os.getenv("USER"), + database="postgres", + timeout=1, + ) + await conn.close() + return + except (asyncpg.CannotConnectNowError, OSError): + await asyncio.sleep(0.5) + + raise RuntimeError( + f"PostgreSQL ({self.instance.name}) did not start within {timeout}s" + ) + + async def _create_pool(self): + """Create connection pool.""" + self.pool = await asyncpg.create_pool( + host=self.instance.host, + port=self.instance.port, + user=self.instance.user or os.getenv("USER"), + database="postgres", + min_size=2, + max_size=10, + ) + + async def stop(self): + """Stop the PostgreSQL instance.""" + if self.process is None: + return + + logger.info("Stopping PostgreSQL (%s)", self.instance.name) + + if self.pool: + await self.pool.close() + self.pool = None + + # Send SIGTERM + self.process.terminate() + try: + self.process.wait(timeout=10) + except subprocess.TimeoutExpired: + logger.warning("PostgreSQL did not stop gracefully, sending SIGKILL") + self.process.kill() + self.process.wait() + + self.process = None + logger.info("PostgreSQL (%s) stopped", self.instance.name) + + async def execute(self, query: str): + """Execute a query.""" + async with self.pool.acquire() as conn: + return await conn.execute(query) + + async def fetch(self, query: str): + """Fetch query results.""" + async with self.pool.acquire() as conn: + return await conn.fetch(query) + + async def fetchrow(self, query: str): + """Fetch single row.""" + async with self.pool.acquire() as conn: + return await conn.fetchrow(query) + + +class AtomicsBenchmarkSuite: + """Main benchmark orchestrator.""" + + def __init__(self, config: Optional[AtomicsBenchmarkConfig] = None): + self.config = config or AtomicsBenchmarkConfig() + + self.traditional_mgr = PostgresManager(self.config.traditional_instance) + self.stdatomic_mgr = PostgresManager(self.config.stdatomic_instance) + + self.data_generator = DataGenerator(seed=self.config.seed) + self.results: List[ComparisonResult] = [] + + async def setup(self): + """Initialize both PostgreSQL instances.""" + logger.info("Setting up benchmark environment...") + + # Initialize clusters + await self.traditional_mgr.initialize_cluster() + await self.stdatomic_mgr.initialize_cluster() + + # Start instances + await self.traditional_mgr.start() + await self.stdatomic_mgr.start() + + # Create benchmark database + await self.traditional_mgr.execute("CREATE DATABASE benchmark_db") + await self.stdatomic_mgr.execute("CREATE DATABASE benchmark_db") + + logger.info("Benchmark environment ready") + + async def teardown(self): + """Stop both PostgreSQL instances.""" + logger.info("Tearing down benchmark environment...") + await self.traditional_mgr.stop() + await self.stdatomic_mgr.stop() + + async def run_benchmark( + self, + schema: TableSchema, + row_count: int, + distribution: DataDistribution, + ) -> List[ComparisonResult]: + """Run complete benchmark for one scenario, comparing both implementations.""" + logger.info( + "=== Benchmark: %s, %d rows, %s ===", + schema.name, + row_count, + distribution.value, + ) + + results = [] + + # Run on traditional + trad_workload, trad_metrics = await self._run_on_instance( + self.traditional_mgr, "traditional", schema, row_count, distribution + ) + + # Run on stdatomic + std_workload, std_metrics = await self._run_on_instance( + self.stdatomic_mgr, "stdatomic", schema, row_count, distribution + ) + + # Compare results + for trad_result in trad_workload.results: + pattern = trad_result.query_pattern + std_result = next( + (r for r in std_workload.results if r.query_pattern == pattern), None + ) + if std_result is None: + continue + + speedup = trad_result.elapsed_seconds / std_result.elapsed_seconds + + comparison = ComparisonResult( + schema_name=schema.name, + row_count=row_count, + distribution=distribution.value, + pattern=pattern, + traditional_ms=trad_result.elapsed_seconds * 1000, + stdatomic_ms=std_result.elapsed_seconds * 1000, + speedup=speedup, + traditional_metrics=trad_metrics, + stdatomic_metrics=std_metrics, + ) + results.append(comparison) + + return results + + async def _run_on_instance( + self, + mgr: PostgresManager, + impl_name: str, + schema: TableSchema, + row_count: int, + distribution: DataDistribution, + ) -> Tuple[WorkloadResult, BenchmarkMetrics]: + """Run benchmark on a single PostgreSQL instance.""" + # Create table + table_name = f"{schema.name}_{distribution.value}" + await self._create_table(mgr, table_name, schema) + + # Load data + load_start = time.perf_counter() + insert_sql = self.data_generator.generate_server_side_insert( + schema, row_count, distribution, table_suffix=f"_{distribution.value}" + ) + await mgr.execute(insert_sql) + load_time = time.perf_counter() - load_start + + logger.info("%s: loaded %d rows in %.2fs", impl_name, row_count, load_time) + + # Collect pre-workload metrics + metrics = await self._collect_metrics( + mgr, impl_name, table_name, schema.name, row_count, distribution.value + ) + metrics.load_time_seconds = load_time + + # Run workload + runner = AtomicsWorkloadRunner( + mgr.pool, + warmup_iterations=self.config.warmup_iterations, + measure_iterations=self.config.measure_iterations, + ) + + workload_result = await runner.run_workload( + schema=schema, + table_name=table_name, + row_count=row_count, + distribution=distribution.value, + patterns=self.config.query_patterns, + ) + + # Cleanup + await mgr.execute(f"DROP TABLE IF EXISTS {table_name} CASCADE") + + return workload_result, metrics + + async def _create_table( + self, mgr: PostgresManager, table_name: str, schema: TableSchema + ): + """Create a table with the given schema.""" + col_defs = [] + for col_name, col_type in schema.columns: + col_defs.append(f"{col_name} {col_type.value}") + + create_sql = f"CREATE TABLE {table_name} ({', '.join(col_defs)})" + await mgr.execute(create_sql) + + # Create indexes + for idx_col in schema.index_columns: + idx_name = f"{table_name}_{idx_col}_idx" + await mgr.execute( + f"CREATE INDEX {idx_name} ON {table_name} ({idx_col})" + ) + + async def _collect_metrics( + self, + mgr: PostgresManager, + impl_name: str, + table_name: str, + schema_name: str, + row_count: int, + distribution: str, + ) -> BenchmarkMetrics: + """Collect metrics for a table.""" + # Table size + size_row = await mgr.fetchrow( + f"SELECT pg_total_relation_size('{table_name}') as total_size" + ) + table_size = size_row["total_size"] + + # Index size + idx_size_row = await mgr.fetchrow( + f"SELECT pg_indexes_size('{table_name}') as idx_size" + ) + index_size = idx_size_row["idx_size"] + + return BenchmarkMetrics( + implementation=impl_name, + schema_name=schema_name, + row_count=row_count, + distribution=distribution, + load_time_seconds=0.0, # Set by caller + table_size_bytes=table_size, + index_size_bytes=index_size, + ) + + async def run_full_suite(self) -> Dict: + """Run complete benchmark suite.""" + start_time = time.perf_counter() + all_results = [] + + for schema in self.config.schemas: + for row_count in self.config.get_row_counts(): + for distribution in self.config.distributions: + try: + results = await self.run_benchmark( + schema, row_count, distribution + ) + all_results.extend(results) + except Exception as e: + logger.error( + "Benchmark failed for %s/%d/%s: %s", + schema.name, + row_count, + distribution.value, + e, + exc_info=True, + ) + + elapsed_time = time.perf_counter() - start_time + + # Generate summary + summary = self._generate_summary(all_results, elapsed_time) + self.results = all_results + + return summary + + def _generate_summary( + self, results: List[ComparisonResult], elapsed_time: float + ) -> Dict: + """Generate summary statistics.""" + if not results: + return {} + + speedups = [r.speedup for r in results] + + summary = { + "total_scenarios": len(results), + "elapsed_time_seconds": elapsed_time, + "median_speedup": sorted(speedups)[len(speedups) // 2], + "mean_speedup": sum(speedups) / len(speedups), + "min_speedup": min(speedups), + "max_speedup": max(speedups), + "stdatomic_faster_count": sum(1 for s in speedups if s > 1.0), + "traditional_faster_count": sum(1 for s in speedups if s < 1.0), + "neutral_count": sum(1 for s in speedups if s == 1.0), + } + + # Per-pattern breakdown + pattern_speedups = {} + for result in results: + if result.pattern not in pattern_speedups: + pattern_speedups[result.pattern] = [] + pattern_speedups[result.pattern].append(result.speedup) + + summary["per_pattern_avg_speedup"] = { + pattern: sum(speedups) / len(speedups) + for pattern, speedups in pattern_speedups.items() + } + + # Find best/worst scenarios + best = max(results, key=lambda r: r.speedup) + worst = min(results, key=lambda r: r.speedup) + + summary["best_stdatomic_scenario"] = { + "pattern": best.pattern, + "schema": best.schema_name, + "distribution": best.distribution, + "speedup": best.speedup, + } + + summary["worst_stdatomic_scenario"] = { + "pattern": worst.pattern, + "schema": worst.schema_name, + "distribution": worst.distribution, + "speedup": worst.speedup, + } + + return summary + + +async def run_atomics_benchmark( + config: Optional[AtomicsBenchmarkConfig] = None, +) -> Dict: + """Main entry point for running atomics benchmark.""" + suite = AtomicsBenchmarkSuite(config) + + try: + await suite.setup() + summary = await suite.run_full_suite() + return {"summary": summary, "results": suite.results} + finally: + await suite.teardown() diff --git a/src/test/benchmarks/atomics_cli.py b/src/test/benchmarks/atomics_cli.py new file mode 100644 index 00000000000..20d870dc6b7 --- /dev/null +++ b/src/test/benchmarks/atomics_cli.py @@ -0,0 +1,354 @@ +""" +CLI entry point for atomics benchmark suite. + +Usage: + python -m src.test.benchmarks.atomics_cli [OPTIONS] + +Examples: + # Quick run with defaults + python -m src.test.benchmarks.atomics_cli + + # Full matrix + python -m src.test.benchmarks.atomics_cli --full-matrix + + # Specific schema and row count + python -m src.test.benchmarks.atomics_cli --schema medium --rows 100000 + + # Custom build directories + python -m src.test.benchmarks.atomics_cli \\ + --traditional-build ./build-traditional \\ + --stdatomic-build ./build-stdatomic + + # Verbose output + python -m src.test.benchmarks.atomics_cli -v +""" + +import argparse +import asyncio +import json +import logging +import sys +from pathlib import Path + +from .atomics_config import ( + ALL_SCHEMAS, + AtomicsBenchmarkConfig, + DataDistribution, + MEDIUM_SCHEMA, + NARROW_SCHEMA, + PostgresInstance, + QueryPattern, + WIDE_SCHEMA, +) +from .atomics_benchmark import run_atomics_benchmark + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="PostgreSQL Atomics Performance Benchmark Suite", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + + # Build directories + parser.add_argument( + "--traditional-build", + type=Path, + default=Path.cwd() / "build-traditional", + help="Path to traditional atomics build directory", + ) + parser.add_argument( + "--stdatomic-build", + type=Path, + default=Path.cwd() / "build-stdatomic", + help="Path to stdatomic.h build directory", + ) + + # PostgreSQL instance configuration + parser.add_argument( + "--traditional-port", type=int, default=5433, help="Port for traditional instance" + ) + parser.add_argument( + "--stdatomic-port", type=int, default=5434, help="Port for stdatomic instance" + ) + parser.add_argument( + "--traditional-datadir", + type=Path, + default=Path("/tmp/pgdata-traditional"), + help="Data directory for traditional instance", + ) + parser.add_argument( + "--stdatomic-datadir", + type=Path, + default=Path("/tmp/pgdata-stdatomic"), + help="Data directory for stdatomic instance", + ) + + # Test matrix + parser.add_argument( + "--schema", + choices=["narrow", "medium", "wide", "all"], + default="all", + help="Table schema to test (default: all)", + ) + parser.add_argument( + "--rows", + type=int, + nargs="+", + default=None, + help="Row counts to test (default: 10000 100000 1000000)", + ) + parser.add_argument( + "--distribution", + choices=["random", "clustered", "low_cardinality", "high_null", "all"], + default="all", + help="Data distribution (default: all)", + ) + parser.add_argument( + "--pattern", + choices=[p.value for p in QueryPattern] + ["all"], + default="all", + help="Query pattern to test (default: all)", + ) + parser.add_argument( + "--full-matrix", + action="store_true", + help="Run full matrix including large row counts", + ) + + # Execution + parser.add_argument( + "--warmup", type=int, default=3, help="Warmup iterations (default: 3)" + ) + parser.add_argument( + "--iterations", type=int, default=10, help="Measurement iterations (default: 10)" + ) + parser.add_argument( + "--concurrent-clients", + type=int, + default=8, + help="Number of concurrent clients for write tests (default: 8)", + ) + parser.add_argument("--seed", type=int, default=42, help="RNG seed (default: 42)") + + # Output + parser.add_argument( + "--output-dir", + "-o", + default="atomics_benchmark_results", + help="Output directory", + ) + parser.add_argument( + "-v", "--verbose", action="store_true", help="Verbose logging" + ) + + return parser.parse_args() + + +def build_config(args: argparse.Namespace) -> AtomicsBenchmarkConfig: + """Build configuration from command-line arguments.""" + traditional_inst = PostgresInstance( + name="traditional", + build_dir=args.traditional_build, + port=args.traditional_port, + data_dir=args.traditional_datadir, + ) + + stdatomic_inst = PostgresInstance( + name="stdatomic", + build_dir=args.stdatomic_build, + port=args.stdatomic_port, + data_dir=args.stdatomic_datadir, + ) + + schema_map = { + "narrow": [NARROW_SCHEMA], + "medium": [MEDIUM_SCHEMA], + "wide": [WIDE_SCHEMA], + "all": list(ALL_SCHEMAS), + } + schemas = schema_map[args.schema] + + if args.distribution == "all": + distributions = [ + DataDistribution.RANDOM, + DataDistribution.CLUSTERED, + DataDistribution.LOW_CARDINALITY, + ] + else: + distributions = [DataDistribution(args.distribution)] + + if args.pattern == "all": + # Default patterns (exclude concurrent patterns for now) + patterns = [ + QueryPattern.FULL_SCAN, + QueryPattern.COLUMN_PROJECTION, + QueryPattern.FILTERED_SCAN, + QueryPattern.AGGREGATION, + QueryPattern.GROUP_BY, + QueryPattern.INDEX_SCAN, + ] + else: + patterns = [QueryPattern(args.pattern)] + + config = AtomicsBenchmarkConfig( + traditional_instance=traditional_inst, + stdatomic_instance=stdatomic_inst, + schemas=schemas, + distributions=distributions, + query_patterns=patterns, + warmup_iterations=args.warmup, + measure_iterations=args.iterations, + concurrent_clients=args.concurrent_clients, + seed=args.seed, + output_dir=args.output_dir, + full_matrix=args.full_matrix, + verbose=args.verbose, + ) + + if args.rows: + config.row_counts = args.rows + + return config + + +def print_summary(summary: dict): + """Print benchmark summary.""" + print() + print("=" * 70) + print(" ATOMICS BENCHMARK RESULTS SUMMARY") + print("=" * 70) + + if not summary: + print(" No results to display") + return + + print(f" Total scenarios tested: {summary.get('total_scenarios', 0)}") + print(f" Total time: {summary.get('elapsed_time_seconds', 0):.1f}s") + print() + + median = summary.get("median_speedup", 1.0) + mean = summary.get("mean_speedup", 1.0) + print(f" Median speedup (std/trad): {median:.3f}x") + print(f" Mean speedup: {mean:.3f}x") + print(f" Min speedup: {summary.get('min_speedup', 1.0):.3f}x") + print(f" Max speedup: {summary.get('max_speedup', 1.0):.3f}x") + print() + + faster = summary.get("stdatomic_faster_count", 0) + slower = summary.get("traditional_faster_count", 0) + total = summary.get("total_scenarios", 1) + + print(f" Stdatomic faster: {faster} ({100*faster/total:.1f}%)") + print(f" Traditional faster: {slower} ({100*slower/total:.1f}%)") + + if summary.get("per_pattern_avg_speedup"): + print() + print(" Per-pattern average speedup (stdatomic vs traditional):") + for pattern, speedup in sorted(summary["per_pattern_avg_speedup"].items()): + indicator = ">>>" if speedup > 1.05 else " " + indicator = "<<<" if speedup < 0.95 else indicator + print(f" {indicator} {pattern:30s} {speedup:.3f}x") + + if summary.get("best_stdatomic_scenario"): + best = summary["best_stdatomic_scenario"] + print() + print( + f" Best stdatomic: {best['pattern']} on {best['schema']} " + f"({best['distribution']}) = {best['speedup']:.3f}x" + ) + + if summary.get("worst_stdatomic_scenario"): + worst = summary["worst_stdatomic_scenario"] + print( + f" Worst stdatomic: {worst['pattern']} on {worst['schema']} " + f"({worst['distribution']}) = {worst['speedup']:.3f}x" + ) + + print("=" * 70) + + +def main(): + args = parse_args() + + log_level = logging.DEBUG if args.verbose else logging.INFO + logging.basicConfig( + level=log_level, + format="%(asctime)s %(levelname)-8s %(name)s: %(message)s", + datefmt="%H:%M:%S", + ) + + # Verify build directories exist + if not args.traditional_build.exists(): + print(f"ERROR: Traditional build directory not found: {args.traditional_build}") + print("Please build with: meson setup build-traditional -Duse_stdatomic=no") + sys.exit(1) + + if not args.stdatomic_build.exists(): + print(f"ERROR: Stdatomic build directory not found: {args.stdatomic_build}") + print("Please build with: meson setup build-stdatomic -Duse_stdatomic=yes") + sys.exit(1) + + config = build_config(args) + + print("=" * 70) + print(" PostgreSQL Atomics Performance Benchmark") + print("=" * 70) + print(f" Traditional build: {config.traditional_instance.build_dir}") + print(f" Stdatomic build: {config.stdatomic_instance.build_dir}") + print(f" Schemas: {[s.name for s in config.schemas]}") + print(f" Row counts: {config.get_row_counts()}") + print(f" Distributions: {[d.value for d in config.distributions]}") + print(f" Patterns: {[p.value for p in config.query_patterns]}") + print( + f" Iterations: {config.measure_iterations} " + f"(warmup: {config.warmup_iterations})" + ) + print(f" Output: {config.output_dir}") + print("=" * 70) + print() + + try: + report = asyncio.run(run_atomics_benchmark(config)) + except KeyboardInterrupt: + print("\nBenchmark interrupted.") + sys.exit(1) + except Exception as e: + logging.error("Benchmark failed: %s", e, exc_info=True) + sys.exit(1) + + # Print summary + print_summary(report.get("summary", {})) + + # Save results + output_dir = Path(config.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + results_file = output_dir / "results.json" + with open(results_file, "w") as f: + # Convert results to serializable format + results_data = [] + for r in report.get("results", []): + results_data.append( + { + "schema": r.schema_name, + "row_count": r.row_count, + "distribution": r.distribution, + "pattern": r.pattern, + "traditional_ms": r.traditional_ms, + "stdatomic_ms": r.stdatomic_ms, + "speedup": r.speedup, + } + ) + + json.dump( + {"summary": report.get("summary", {}), "results": results_data}, + f, + indent=2, + ) + + print(f"\nResults saved to: {results_file}") + + +if __name__ == "__main__": + main() diff --git a/src/test/benchmarks/atomics_config.py b/src/test/benchmarks/atomics_config.py new file mode 100644 index 00000000000..0b2c831cbf2 --- /dev/null +++ b/src/test/benchmarks/atomics_config.py @@ -0,0 +1,275 @@ +""" +Benchmark configuration for comparing traditional vs stdatomic.h implementations. + +This module defines the test matrix for evaluating atomics performance impact +across different workload patterns, table schemas, and data distributions. +""" + +import os +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import List, Optional + + +class TableWidth(Enum): + """Table width categories for testing atomic operation overhead.""" + NARROW = "narrow" # 3-5 columns - minimal overhead + MEDIUM = "medium" # 10-30 columns - moderate overhead + WIDE = "wide" # 50-120 columns - heavy overhead + + +class DataDistribution(Enum): + """Data distribution patterns affecting atomic contention.""" + RANDOM = "random" + CLUSTERED = "clustered" + LOW_CARDINALITY = "low_cardinality" + HIGH_NULL = "high_null" + + +class QueryPattern(Enum): + """Query patterns that stress different atomic operations.""" + FULL_SCAN = "full_scan" # Sequential scan - spinlock acquisition + COLUMN_PROJECTION = "column_projection" # Selective scan - tuple visibility checks + FILTERED_SCAN = "filtered_scan" # Index + seq scan - mixed contention + AGGREGATION = "aggregation" # Heavy computation - shared buffer access + GROUP_BY = "group_by" # Hash aggregation - memory ordering stress + INDEX_SCAN = "index_scan" # Index-only - buffer pin/unpin contention + CONCURRENT_INSERT = "concurrent_insert" # High contention writes + CONCURRENT_UPDATE = "concurrent_update" # Lock management stress + MIXED_WORKLOAD = "mixed_workload" # Read/write mix - realistic contention + + +class ColumnType(Enum): + """Column types for schema generation.""" + INT = "integer" + BIGINT = "bigint" + TEXT = "text" + BOOLEAN = "boolean" + UUID = "uuid" + TIMESTAMP = "timestamp" + FLOAT = "double precision" + NUMERIC = "numeric(12,2)" + JSONB = "jsonb" + + +# Row counts for testing atomic overhead at different scales +ROW_COUNTS = [1_000, 10_000, 100_000, 1_000_000, 10_000_000] + +# Default subset for quick runs +DEFAULT_ROW_COUNTS = [10_000, 100_000, 1_000_000] + + +@dataclass +class PostgresInstance: + """Configuration for a PostgreSQL instance (traditional or stdatomic build).""" + name: str # "traditional" or "stdatomic" + build_dir: Path # Path to build or install directory + host: str = "localhost" + port: int = 5432 + data_dir: Optional[Path] = None + user: str = "" + password: str = "" + + def _is_install_dir(self) -> bool: + """Check if build_dir is an install directory (has bin/ subdirectory).""" + return (self.build_dir / "bin").exists() + + @property + def postgres_bin(self) -> Path: + """Path to postgres binary.""" + if self._is_install_dir(): + return self.build_dir / "bin" / "postgres" + return self.build_dir / "src" / "backend" / "postgres" + + @property + def initdb_bin(self) -> Path: + """Path to initdb binary.""" + if self._is_install_dir(): + return self.build_dir / "bin" / "initdb" + return self.build_dir / "src" / "bin" / "initdb" / "initdb" + + @property + def pg_ctl_bin(self) -> Path: + """Path to pg_ctl binary.""" + if self._is_install_dir(): + return self.build_dir / "bin" / "pg_ctl" + return self.build_dir / "src" / "bin" / "pg_ctl" / "pg_ctl" + + def get_dsn(self, database: str = "postgres") -> str: + """Get connection DSN.""" + parts = [ + f"host={self.host}", + f"port={self.port}", + f"dbname={database}", + ] + if self.user: + parts.append(f"user={self.user}") + if self.password: + parts.append(f"password={self.password}") + return " ".join(parts) + + +@dataclass +class TableSchema: + """Defines a table schema for benchmarking.""" + name: str + width: TableWidth + columns: List[tuple] # (col_name, ColumnType) + index_columns: List[str] = field(default_factory=list) + + @property + def column_names(self) -> List[str]: + return [c[0] for c in self.columns] + + @property + def column_types(self) -> List[ColumnType]: + return [c[1] for c in self.columns] + + +# Pre-defined table schemas for the test matrix +NARROW_SCHEMA = TableSchema( + name="atomics_narrow", + width=TableWidth.NARROW, + columns=[ + ("id", ColumnType.BIGINT), + ("val_int", ColumnType.INT), + ("val_text", ColumnType.TEXT), + ("flag", ColumnType.BOOLEAN), + ], + index_columns=["id"], +) + +MEDIUM_SCHEMA = TableSchema( + name="atomics_medium", + width=TableWidth.MEDIUM, + columns=[ + ("id", ColumnType.BIGINT), + ("category", ColumnType.INT), + ("amount", ColumnType.NUMERIC), + ("description", ColumnType.TEXT), + ("is_active", ColumnType.BOOLEAN), + ("created_at", ColumnType.TIMESTAMP), + ("ref_uuid", ColumnType.UUID), + ("score", ColumnType.FLOAT), + ("status_code", ColumnType.INT), + ("notes", ColumnType.TEXT), + ("metadata", ColumnType.JSONB), + ], + index_columns=["id", "category"], +) + + +def _build_wide_columns(): + """Build a wide schema with 55 columns covering all data types.""" + cols = [("id", ColumnType.BIGINT)] + # 8 INT columns + for i in range(1, 9): + cols.append((f"col_int_{i}", ColumnType.INT)) + # 5 BIGINT columns + for i in range(1, 6): + cols.append((f"col_bigint_{i}", ColumnType.BIGINT)) + # 8 TEXT columns + for i in range(1, 9): + cols.append((f"col_text_{i}", ColumnType.TEXT)) + # 6 BOOLEAN columns + for i in range(1, 7): + cols.append((f"col_bool_{i}", ColumnType.BOOLEAN)) + # 5 FLOAT columns + for i in range(1, 6): + cols.append((f"col_float_{i}", ColumnType.FLOAT)) + # 5 NUMERIC columns + for i in range(1, 6): + cols.append((f"col_numeric_{i}", ColumnType.NUMERIC)) + # 5 UUID columns + for i in range(1, 6): + cols.append((f"col_uuid_{i}", ColumnType.UUID)) + # 5 TIMESTAMP columns + for i in range(1, 6): + cols.append((f"col_ts_{i}", ColumnType.TIMESTAMP)) + # 4 JSONB columns + for i in range(1, 5): + cols.append((f"col_jsonb_{i}", ColumnType.JSONB)) + # 3 more INT columns to reach 55 + for i in range(9, 12): + cols.append((f"col_int_{i}", ColumnType.INT)) + return cols + + +WIDE_SCHEMA = TableSchema( + name="atomics_wide", + width=TableWidth.WIDE, + columns=_build_wide_columns(), + index_columns=["id", "col_int_1", "col_text_1"], +) + +ALL_SCHEMAS = [NARROW_SCHEMA, MEDIUM_SCHEMA, WIDE_SCHEMA] + + +@dataclass +class AtomicsBenchmarkConfig: + """Configuration for atomics performance comparison.""" + # PostgreSQL instances to compare + traditional_instance: Optional[PostgresInstance] = None + stdatomic_instance: Optional[PostgresInstance] = None + + # Test matrix + schemas: List[TableSchema] = field(default_factory=lambda: list(ALL_SCHEMAS)) + row_counts: List[int] = field(default_factory=lambda: list(DEFAULT_ROW_COUNTS)) + distributions: List[DataDistribution] = field( + default_factory=lambda: [ + DataDistribution.RANDOM, + DataDistribution.CLUSTERED, + DataDistribution.LOW_CARDINALITY, + ] + ) + query_patterns: List[QueryPattern] = field( + default_factory=lambda: [ + QueryPattern.FULL_SCAN, + QueryPattern.COLUMN_PROJECTION, + QueryPattern.FILTERED_SCAN, + QueryPattern.AGGREGATION, + QueryPattern.GROUP_BY, + QueryPattern.INDEX_SCAN, + ] + ) + + # Execution parameters + warmup_iterations: int = 3 + measure_iterations: int = 10 + concurrent_clients: int = 8 # For concurrent workload tests + seed: int = 42 + + # Output + output_dir: str = "atomics_benchmark_results" + verbose: bool = False + + # Run the full matrix or a reduced subset + full_matrix: bool = False + + def __post_init__(self): + """Initialize default instances if not provided.""" + if self.traditional_instance is None: + build_dir = Path.cwd() / "build-traditional" + if not build_dir.exists(): + build_dir = Path.cwd() / "build" + self.traditional_instance = PostgresInstance( + name="traditional", + build_dir=build_dir, + port=5433, + data_dir=Path("/tmp/pgdata-traditional"), + ) + + if self.stdatomic_instance is None: + build_dir = Path.cwd() / "build-stdatomic" + self.stdatomic_instance = PostgresInstance( + name="stdatomic", + build_dir=build_dir, + port=5434, + data_dir=Path("/tmp/pgdata-stdatomic"), + ) + + def get_row_counts(self) -> List[int]: + if self.full_matrix: + return ROW_COUNTS + return self.row_counts diff --git a/src/test/benchmarks/atomics_workload.py b/src/test/benchmarks/atomics_workload.py new file mode 100644 index 00000000000..0a64e974535 --- /dev/null +++ b/src/test/benchmarks/atomics_workload.py @@ -0,0 +1,304 @@ +""" +Workload runner for atomics benchmarking. + +Executes query patterns that stress different atomic operations: +- Buffer management (spinlocks) +- Tuple visibility (atomic reads) +- Lock management (test-and-set) +- Shared memory access (memory ordering) +""" + +import asyncio +import asyncpg +import logging +import time +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from .atomics_config import ColumnType, QueryPattern, TableSchema + +logger = logging.getLogger(__name__) + + +@dataclass +class QueryResult: + """Result of a single query execution.""" + query_pattern: str + table_name: str + implementation: str # "traditional" or "stdatomic" + query_sql: str + elapsed_seconds: float + row_count: int = 0 + explain_plan: Optional[Dict[str, Any]] = None + + +@dataclass +class WorkloadResult: + """Aggregated results for a complete workload run.""" + schema_name: str + row_count: int + distribution: str + implementation: str + results: List[QueryResult] = field(default_factory=list) + + def add(self, result: QueryResult): + self.results.append(result) + + +class AtomicsWorkloadRunner: + """Generates and executes query workloads for atomics testing.""" + + def __init__( + self, + pool: asyncpg.Pool, + warmup_iterations: int = 3, + measure_iterations: int = 10, + ): + self.pool = pool + self.warmup_iterations = warmup_iterations + self.measure_iterations = measure_iterations + + # ------------------------------------------------------------------ + # Query generators per pattern + # ------------------------------------------------------------------ + + def _full_scan_query(self, table_name: str, schema: TableSchema) -> str: + """Full table scan - stresses spinlock acquisition and buffer management.""" + return f"SELECT * FROM {table_name}" + + def _column_projection_query(self, table_name: str, schema: TableSchema) -> str: + """Column projection - tuple visibility checks with atomic reads.""" + cols = [c[0] for c in schema.columns if c[0] != "id"][:2] + if not cols: + cols = [schema.columns[0][0]] + return f"SELECT {', '.join(cols)} FROM {table_name}" + + def _filtered_scan_query(self, table_name: str, schema: TableSchema) -> str: + """Filtered scan - mixed index and sequential access patterns.""" + for col_name, col_type in schema.columns: + if col_type == ColumnType.INT and col_name != "id": + return f"SELECT * FROM {table_name} WHERE {col_name} > 0" + if col_type == ColumnType.BOOLEAN: + return f"SELECT * FROM {table_name} WHERE {col_name} = TRUE" + return f"SELECT * FROM {table_name} WHERE id > 100 AND id <= 10000" + + def _aggregation_query(self, table_name: str, schema: TableSchema) -> str: + """Aggregation - heavy shared buffer access with memory ordering stress.""" + agg_exprs = [] + for col_name, col_type in schema.columns: + if col_type in ( + ColumnType.INT, + ColumnType.BIGINT, + ColumnType.FLOAT, + ColumnType.NUMERIC, + ): + agg_exprs.append(f"SUM({col_name})") + agg_exprs.append(f"AVG({col_name})") + if len(agg_exprs) >= 6: + break + if not agg_exprs: + agg_exprs = ["COUNT(*)"] + return f"SELECT COUNT(*), {', '.join(agg_exprs)} FROM {table_name}" + + def _group_by_query(self, table_name: str, schema: TableSchema) -> str: + """Group-by with hash aggregation - memory ordering and buffer management.""" + group_col = None + agg_col = None + for col_name, col_type in schema.columns: + if col_name == "id": + continue + if col_type in (ColumnType.INT, ColumnType.BOOLEAN) and group_col is None: + group_col = col_name + if ( + col_type + in (ColumnType.FLOAT, ColumnType.NUMERIC, ColumnType.INT, ColumnType.BIGINT) + and agg_col is None + ): + agg_col = col_name + + if group_col is None: + group_col = schema.columns[0][0] + if agg_col is None: + agg_col = "id" + + return ( + f"SELECT {group_col}, COUNT(*), SUM({agg_col}), AVG({agg_col}) " + f"FROM {table_name} GROUP BY {group_col}" + ) + + def _index_scan_query(self, table_name: str, schema: TableSchema) -> str: + """Index scan - buffer pin/unpin contention.""" + return f"SELECT * FROM {table_name} WHERE id BETWEEN 100 AND 200" + + def _concurrent_insert_query(self, table_name: str, schema: TableSchema) -> str: + """Concurrent inserts - high contention lock management.""" + # Generate a simple insert + col_names = schema.column_names + values = [] + for col_name, col_type in schema.columns: + if col_name == "id": + values.append("nextval('seq_concurrent_insert')") + elif col_type == ColumnType.INT: + values.append("42") + elif col_type == ColumnType.BIGINT: + values.append("12345") + elif col_type == ColumnType.TEXT: + values.append("'test'") + elif col_type == ColumnType.BOOLEAN: + values.append("TRUE") + elif col_type == ColumnType.UUID: + values.append("gen_random_uuid()") + elif col_type == ColumnType.TIMESTAMP: + values.append("NOW()") + elif col_type == ColumnType.FLOAT: + values.append("3.14") + elif col_type == ColumnType.NUMERIC: + values.append("99.99") + elif col_type == ColumnType.JSONB: + values.append("'{}'::jsonb") + else: + values.append("NULL") + + return f"INSERT INTO {table_name} ({', '.join(col_names)}) VALUES ({', '.join(values)})" + + def _concurrent_update_query(self, table_name: str, schema: TableSchema) -> str: + """Concurrent updates - lock management and visibility checks.""" + update_col = None + for col_name, col_type in schema.columns: + if col_name != "id" and col_type in (ColumnType.INT, ColumnType.BIGINT): + update_col = col_name + break + + if update_col is None: + update_col = "id" + + return f"UPDATE {table_name} SET {update_col} = {update_col} + 1 WHERE id <= 100" + + def _mixed_workload_query(self, table_name: str, schema: TableSchema) -> str: + """Mixed read/write workload - realistic contention patterns.""" + return f""" + WITH updated AS ( + UPDATE {table_name} SET id = id WHERE id <= 10 RETURNING * + ) + SELECT COUNT(*) FROM {table_name} WHERE id > 10 + """ + + def _get_query( + self, pattern: QueryPattern, table_name: str, schema: TableSchema + ) -> str: + """Get query for a given pattern.""" + generators = { + QueryPattern.FULL_SCAN: self._full_scan_query, + QueryPattern.COLUMN_PROJECTION: self._column_projection_query, + QueryPattern.FILTERED_SCAN: self._filtered_scan_query, + QueryPattern.AGGREGATION: self._aggregation_query, + QueryPattern.GROUP_BY: self._group_by_query, + QueryPattern.INDEX_SCAN: self._index_scan_query, + QueryPattern.CONCURRENT_INSERT: self._concurrent_insert_query, + QueryPattern.CONCURRENT_UPDATE: self._concurrent_update_query, + QueryPattern.MIXED_WORKLOAD: self._mixed_workload_query, + } + gen = generators.get(pattern) + if gen is None: + raise ValueError(f"Unknown query pattern: {pattern}") + return gen(table_name, schema) + + # ------------------------------------------------------------------ + # Execution + # ------------------------------------------------------------------ + + async def _run_single( + self, + query: str, + pattern: QueryPattern, + table_name: str, + implementation: str, + collect_explain: bool = False, + ) -> QueryResult: + """Run a single query, returning timing and optional EXPLAIN data.""" + # Warmup + async with self.pool.acquire() as conn: + for _ in range(self.warmup_iterations): + try: + await conn.fetch(query) + except Exception as e: + logger.warning("Warmup query failed: %s", e) + + # Measure + timings = [] + for _ in range(self.measure_iterations): + start = time.perf_counter() + async with self.pool.acquire() as conn: + try: + rows = await conn.fetch(query) + elapsed = time.perf_counter() - start + timings.append(elapsed) + except Exception as e: + logger.error("Measurement query failed: %s", e) + elapsed = 0.0 + timings.append(elapsed) + + # Use median timing + timings.sort() + median_time = timings[len(timings) // 2] + + # Optional EXPLAIN + explain_plan = None + if collect_explain: + try: + async with self.pool.acquire() as conn: + explain_result = await conn.fetch( + f"EXPLAIN (FORMAT JSON, ANALYZE, BUFFERS) {query}" + ) + if explain_result: + explain_plan = explain_result[0][0] + except Exception as e: + logger.warning("EXPLAIN failed: %s", e) + + return QueryResult( + query_pattern=pattern.value, + table_name=table_name, + implementation=implementation, + query_sql=query, + elapsed_seconds=median_time, + row_count=len(rows) if rows else 0, + explain_plan=explain_plan, + ) + + async def run_workload( + self, + schema: TableSchema, + table_name: str, + row_count: int, + distribution: str, + patterns: List[QueryPattern], + ) -> WorkloadResult: + """Run complete workload and return results.""" + result = WorkloadResult( + schema_name=schema.name, + row_count=row_count, + distribution=distribution, + implementation="", # Set by caller + ) + + for pattern in patterns: + try: + query = self._get_query(pattern, table_name, schema) + query_result = await self._run_single( + query=query, + pattern=pattern, + table_name=table_name, + implementation=result.implementation, + collect_explain=False, # Too expensive for full matrix + ) + result.add(query_result) + + logger.info( + " %s: %.2fms", + pattern.value, + query_result.elapsed_seconds * 1000, + ) + except Exception as e: + logger.error("Pattern %s failed: %s", pattern.value, e, exc_info=True) + + return result diff --git a/src/test/benchmarks/data_generator.py b/src/test/benchmarks/data_generator.py new file mode 100644 index 00000000000..2366700c6ee --- /dev/null +++ b/src/test/benchmarks/data_generator.py @@ -0,0 +1,409 @@ +""" +Reproducible seeded random data generation for benchmark tables. + +Generates SQL INSERT statements or COPY-compatible data for various +column types and data distributions. +""" + +import hashlib +import logging +import random +import uuid +from datetime import datetime, timedelta +from typing import Any, List, Optional + +from .atomics_config import ColumnType, DataDistribution, TableSchema + +logger = logging.getLogger(__name__) + +# Low-cardinality value pools +LOW_CARD_TEXT = [ + "active", "inactive", "pending", "completed", "cancelled", + "processing", "shipped", "returned", "refunded", "on_hold", +] +LOW_CARD_INT_RANGE = 20 +LOW_CARD_STATUS_CODES = [100, 200, 201, 301, 400, 403, 404, 500, 502, 503] + +# Clustered parameters +CLUSTER_CENTERS = 5 +CLUSTER_SPREAD = 100 + +# Base timestamp for reproducible timestamp generation +BASE_TS = datetime(2020, 1, 1) + + +class DataGenerator: + """Generates reproducible test data for benchmark tables.""" + + def __init__(self, seed: int = 42): + self.seed = seed + self._rng = random.Random(seed) + + def reset(self): + """Reset the RNG to produce identical sequences.""" + self._rng = random.Random(self.seed) + + # ------------------------------------------------------------------ + # Value generators per column type and distribution + # ------------------------------------------------------------------ + + def _gen_int(self, dist: DataDistribution, row_idx: int) -> int: + if dist == DataDistribution.RANDOM: + return self._rng.randint(-2_147_483_648, 2_147_483_647) + elif dist == DataDistribution.CLUSTERED: + center = (row_idx % CLUSTER_CENTERS) * 1_000_000 + return center + self._rng.randint(-CLUSTER_SPREAD, CLUSTER_SPREAD) + else: # LOW_CARDINALITY + return self._rng.choice(LOW_CARD_STATUS_CODES) + + def _gen_bigint(self, dist: DataDistribution, row_idx: int) -> int: + if dist == DataDistribution.RANDOM: + return self._rng.randint(0, 2**62) + elif dist == DataDistribution.CLUSTERED: + center = (row_idx % CLUSTER_CENTERS) * 10_000_000_000 + return center + self._rng.randint(-1000, 1000) + else: + return self._rng.randint(1, LOW_CARD_INT_RANGE) + + def _gen_text(self, dist: DataDistribution, row_idx: int) -> str: + if dist == DataDistribution.RANDOM: + # MD5-like random string + h = hashlib.md5(f"{self.seed}-{row_idx}-{self._rng.random()}".encode()) + return h.hexdigest() + elif dist == DataDistribution.CLUSTERED: + group = row_idx % CLUSTER_CENTERS + suffix = self._rng.randint(0, CLUSTER_SPREAD) + return f"group_{group}_item_{suffix}" + else: + return self._rng.choice(LOW_CARD_TEXT) + + def _gen_boolean(self, dist: DataDistribution, row_idx: int) -> bool: + if dist == DataDistribution.RANDOM: + return self._rng.random() < 0.5 + elif dist == DataDistribution.CLUSTERED: + # Runs of True/False + return (row_idx // 100) % 2 == 0 + else: + # Heavily skewed: 95% True + return self._rng.random() < 0.95 + + def _gen_uuid(self, dist: DataDistribution, row_idx: int) -> str: + if dist == DataDistribution.LOW_CARDINALITY: + # Only 10 distinct UUIDs + idx = row_idx % 10 + return str(uuid.UUID(int=idx + 1)) + # For RANDOM and CLUSTERED, use seeded generation + bits = self._rng.getrandbits(128) + return str(uuid.UUID(int=bits, version=4)) + + def _gen_timestamp(self, dist: DataDistribution, row_idx: int) -> str: + if dist == DataDistribution.RANDOM: + days = self._rng.randint(0, 1825) # ~5 years + secs = self._rng.randint(0, 86400) + ts = BASE_TS + timedelta(days=days, seconds=secs) + elif dist == DataDistribution.CLUSTERED: + # Clustered around specific dates + center_day = (row_idx % CLUSTER_CENTERS) * 365 + offset = self._rng.randint(-30, 30) + ts = BASE_TS + timedelta(days=center_day + offset) + else: + # Low cardinality: 10 distinct dates + day_idx = row_idx % 10 + ts = BASE_TS + timedelta(days=day_idx * 100) + return ts.strftime("%Y-%m-%d %H:%M:%S") + + def _gen_float(self, dist: DataDistribution, row_idx: int) -> float: + if dist == DataDistribution.RANDOM: + return self._rng.uniform(-1e6, 1e6) + elif dist == DataDistribution.CLUSTERED: + center = (row_idx % CLUSTER_CENTERS) * 1000.0 + return center + self._rng.gauss(0, 10) + else: + return self._rng.choice([0.0, 1.0, 10.0, 100.0, 1000.0]) + + def _gen_numeric(self, dist: DataDistribution, row_idx: int) -> str: + val = self._gen_float(dist, row_idx) + return f"{val:.2f}" + + def _gen_jsonb(self, dist: DataDistribution, row_idx: int) -> str: + import json + if dist == DataDistribution.RANDOM: + obj = { + "key": self._rng.randint(1, 100000), + "label": hashlib.md5(f"{self.seed}-json-{row_idx}".encode()).hexdigest()[:8], + "value": round(self._rng.uniform(0, 1000), 2), + "active": self._rng.random() < 0.5, + } + elif dist == DataDistribution.CLUSTERED: + group = row_idx % CLUSTER_CENTERS + obj = { + "group": group, + "label": f"cluster_{group}", + "value": group * 100 + self._rng.randint(0, CLUSTER_SPREAD), + } + elif dist == DataDistribution.HIGH_NULL: + # HIGH_NULL: return None most of the time (handled in _gen_value) + obj = {"id": row_idx % 10, "status": self._rng.choice(LOW_CARD_TEXT)} + else: # LOW_CARDINALITY + obj = {"id": row_idx % 10, "status": self._rng.choice(LOW_CARD_TEXT)} + return json.dumps(obj) + + def _gen_value( + self, col_type: ColumnType, dist: DataDistribution, row_idx: int + ) -> Any: + # HIGH_NULL distribution: ~80% of non-id values are NULL + if dist == DataDistribution.HIGH_NULL and col_type != ColumnType.BIGINT: + if self._rng.random() < 0.80: + return None + + generators = { + ColumnType.INT: self._gen_int, + ColumnType.BIGINT: self._gen_bigint, + ColumnType.TEXT: self._gen_text, + ColumnType.BOOLEAN: self._gen_boolean, + ColumnType.UUID: self._gen_uuid, + ColumnType.TIMESTAMP: self._gen_timestamp, + ColumnType.FLOAT: self._gen_float, + ColumnType.NUMERIC: self._gen_numeric, + ColumnType.JSONB: self._gen_jsonb, + } + gen = generators.get(col_type) + if gen is None: + raise ValueError(f"Unsupported column type: {col_type}") + return gen(dist, row_idx) + + # ------------------------------------------------------------------ + # SQL generation helpers + # ------------------------------------------------------------------ + + def generate_insert_sql( + self, + schema: TableSchema, + row_count: int, + dist: DataDistribution, + table_suffix: str = "", + batch_size: int = 1000, + ) -> List[str]: + """Generate INSERT statements in batches for the given schema. + + Returns a list of SQL strings, each inserting up to batch_size rows. + The ``id`` column is always set to the sequential row index. + """ + self.reset() + col_defs = ", ".join(schema.column_names) + statements = [] + + for batch_start in range(0, row_count, batch_size): + batch_end = min(batch_start + batch_size, row_count) + rows_sql = [] + for i in range(batch_start, batch_end): + vals = [] + for col_name, col_type in schema.columns: + if col_name == "id": + vals.append(str(i + 1)) + else: + v = self._gen_value(col_type, dist, i) + vals.append(self._sql_literal(v, col_type)) + rows_sql.append(f"({', '.join(vals)})") + + table_name = f"{schema.name}{table_suffix}" + stmt = f"INSERT INTO {table_name} ({col_defs}) VALUES\n" + stmt += ",\n".join(rows_sql) + statements.append(stmt) + + return statements + + def generate_copy_data( + self, + schema: TableSchema, + row_count: int, + dist: DataDistribution, + ) -> str: + """Generate tab-separated COPY data for the given schema. + + Returns a single string suitable for COPY ... FROM STDIN. + """ + self.reset() + lines = [] + for i in range(row_count): + vals = [] + for col_name, col_type in schema.columns: + if col_name == "id": + vals.append(str(i + 1)) + else: + v = self._gen_value(col_type, dist, i) + vals.append(self._copy_literal(v, col_type)) + lines.append("\t".join(vals)) + return "\n".join(lines) + + def generate_server_side_insert( + self, + schema: TableSchema, + row_count: int, + dist: DataDistribution, + table_suffix: str = "", + ) -> str: + """Generate a single INSERT ... SELECT generate_series SQL statement. + + This is much faster for large datasets because it runs entirely + server-side without sending row data over the wire. + """ + table_name = f"{schema.name}{table_suffix}" + col_exprs = [] + for col_name, col_type in schema.columns: + if col_name == "id": + col_exprs.append("g AS id") + else: + col_exprs.append( + f"{self._server_side_expr(col_name, col_type, dist, row_count)} AS {col_name}" + ) + + select_list = ",\n ".join(col_exprs) + return ( + f"INSERT INTO {table_name} ({', '.join(schema.column_names)})\n" + f"SELECT {select_list}\n" + f"FROM generate_series(1, {row_count}) AS g" + ) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + @staticmethod + def _sql_literal(value: Any, col_type: ColumnType) -> str: + if value is None: + return "NULL" + if col_type in (ColumnType.TEXT, ColumnType.UUID, ColumnType.TIMESTAMP): + escaped = str(value).replace("'", "''") + return f"'{escaped}'" + if col_type == ColumnType.JSONB: + escaped = str(value).replace("'", "''") + return f"'{escaped}'::jsonb" + if col_type == ColumnType.BOOLEAN: + return "TRUE" if value else "FALSE" + if col_type == ColumnType.NUMERIC: + return str(value) + return str(value) + + @staticmethod + def _copy_literal(value: Any, col_type: ColumnType) -> str: + if value is None: + return "\\N" + if col_type == ColumnType.BOOLEAN: + return "t" if value else "f" + return str(value) + + def _server_side_expr( + self, + col_name: str, + col_type: ColumnType, + dist: DataDistribution, + row_count: int, + ) -> str: + """Return a SQL expression that produces the desired distribution + server-side using generate_series variable ``g``.""" + + seed_val = self.seed + + # HIGH_NULL: wrap the underlying RANDOM expression so ~80% are NULL + if dist == DataDistribution.HIGH_NULL and col_type != ColumnType.BIGINT: + inner = self._server_side_expr( + col_name, col_type, DataDistribution.RANDOM, row_count + ) + return f"CASE WHEN abs(hashint4(g + {seed_val} + 99)) % 5 = 0 THEN {inner} ELSE NULL END" + + if col_type == ColumnType.INT: + if dist == DataDistribution.RANDOM: + return f"(hashint4(g + {seed_val}) % 2147483647)::integer" + elif dist == DataDistribution.CLUSTERED: + return f"((g % {CLUSTER_CENTERS}) * 1000000 + (hashint4(g + {seed_val}) % {CLUSTER_SPREAD}))::integer" + else: + codes = ",".join(str(c) for c in LOW_CARD_STATUS_CODES) + return f"(ARRAY[{codes}])[1 + abs(hashint4(g + {seed_val})) % {len(LOW_CARD_STATUS_CODES)}]" + + if col_type == ColumnType.BIGINT: + if dist == DataDistribution.RANDOM: + return f"(hashint8(g::bigint + {seed_val}) & x'3FFFFFFFFFFFFFFF'::bigint)::bigint" + elif dist == DataDistribution.CLUSTERED: + return f"((g % {CLUSTER_CENTERS})::bigint * 10000000000 + (hashint4(g + {seed_val}) % 1000)::bigint)" + else: + return f"(1 + abs(hashint4(g + {seed_val})) % {LOW_CARD_INT_RANGE})::bigint" + + if col_type == ColumnType.TEXT: + if dist == DataDistribution.RANDOM: + return f"md5(g::text || '{seed_val}')" + elif dist == DataDistribution.CLUSTERED: + return f"'group_' || (g % {CLUSTER_CENTERS})::text || '_item_' || (abs(hashint4(g + {seed_val})) % {CLUSTER_SPREAD})::text" + else: + texts = ",".join(f"'{t}'" for t in LOW_CARD_TEXT) + return f"(ARRAY[{texts}])[1 + abs(hashint4(g + {seed_val})) % {len(LOW_CARD_TEXT)}]" + + if col_type == ColumnType.BOOLEAN: + if dist == DataDistribution.RANDOM: + return f"(hashint4(g + {seed_val}) % 2 = 0)" + elif dist == DataDistribution.CLUSTERED: + return f"((g / 100) % 2 = 0)" + else: + return f"(abs(hashint4(g + {seed_val})) % 20 != 0)" + + if col_type == ColumnType.UUID: + if dist == DataDistribution.LOW_CARDINALITY: + return f"(lpad(((g % 10) + 1)::text, 32, '0'))::uuid" + return f"md5(g::text || '{seed_val}' || random()::text)::uuid" + + if col_type == ColumnType.TIMESTAMP: + if dist == DataDistribution.RANDOM: + return f"'2020-01-01'::timestamp + (abs(hashint4(g + {seed_val})) % 157680000) * interval '1 second'" + elif dist == DataDistribution.CLUSTERED: + return f"'2020-01-01'::timestamp + ((g % {CLUSTER_CENTERS}) * 365 + (abs(hashint4(g + {seed_val})) % 60) - 30) * interval '1 day'" + else: + return f"'2020-01-01'::timestamp + ((g % 10) * 100) * interval '1 day'" + + if col_type == ColumnType.FLOAT: + if dist == DataDistribution.RANDOM: + return f"(hashint4(g + {seed_val})::double precision / 2147483647.0 * 2000000 - 1000000)" + elif dist == DataDistribution.CLUSTERED: + return f"((g % {CLUSTER_CENTERS}) * 1000.0 + (hashint4(g + {seed_val}) % 100)::double precision / 10.0)" + else: + return f"(ARRAY[0.0, 1.0, 10.0, 100.0, 1000.0])[1 + abs(hashint4(g + {seed_val})) % 5]" + + if col_type == ColumnType.NUMERIC: + if dist == DataDistribution.RANDOM: + return f"round((hashint4(g + {seed_val})::numeric / 2147483647.0 * 2000000 - 1000000), 2)" + elif dist == DataDistribution.CLUSTERED: + return f"round(((g % {CLUSTER_CENTERS}) * 1000.0 + (hashint4(g + {seed_val}) % 100)::numeric / 10.0), 2)" + else: + return f"(ARRAY[0.00, 1.00, 10.00, 100.00, 1000.00])[1 + abs(hashint4(g + {seed_val})) % 5]::numeric(12,2)" + + if col_type == ColumnType.JSONB: + if dist == DataDistribution.RANDOM: + return ( + f"jsonb_build_object(" + f"'key', abs(hashint4(g + {seed_val})) % 100000, " + f"'label', left(md5(g::text || '{seed_val}'), 8), " + f"'value', round((hashint4(g + {seed_val})::numeric / 2147483647.0 * 1000), 2), " + f"'active', (hashint4(g + {seed_val}) % 2 = 0))" + ) + elif dist == DataDistribution.CLUSTERED: + return ( + f"jsonb_build_object(" + f"'group', g % {CLUSTER_CENTERS}, " + f"'label', 'cluster_' || (g % {CLUSTER_CENTERS})::text, " + f"'value', (g % {CLUSTER_CENTERS}) * 100 + abs(hashint4(g + {seed_val})) % {CLUSTER_SPREAD})" + ) + elif dist == DataDistribution.HIGH_NULL: + return ( + f"CASE WHEN abs(hashint4(g + {seed_val})) % 5 = 0 THEN " + f"jsonb_build_object('id', g % 10, 'status', " + f"(ARRAY[{','.join(repr(t) for t in LOW_CARD_TEXT)}])" + f"[1 + abs(hashint4(g + {seed_val} + 1)) % {len(LOW_CARD_TEXT)}]) " + f"ELSE NULL END" + ) + else: # LOW_CARDINALITY + texts = ",".join(f"'{t}'" for t in LOW_CARD_TEXT) + return ( + f"jsonb_build_object('id', g % 10, 'status', " + f"(ARRAY[{texts}])[1 + abs(hashint4(g + {seed_val})) % {len(LOW_CARD_TEXT)}])" + ) + + raise ValueError(f"Unsupported column type for server-side generation: {col_type}") -- 2.50.1