From 60c987a891072fcb9473b4b346cea115ffd7214b Mon Sep 17 00:00:00 2001 From: Greg Burd Date: Thu, 9 Jul 2026 12:34:56 -0400 Subject: [PATCH v4 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. --- 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..7d51afa1dc6 --- /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..71fb3e98e95 --- /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