truthound 1.0.8__py3-none-any.whl

truthound/benchmark/generators.py

@@ -0,0 +1,706 @@
+"""Data generators for benchmark scenarios.
+
+This module provides extensible data generators that create realistic
+test data for various benchmarking scenarios.
+
+Generator types:
+- TabularDataGenerator: General-purpose tabular data
+- TimeSeriesDataGenerator: Time-series data with patterns
+- FinancialDataGenerator: Stock/crypto market data
+- TextDataGenerator: Text-heavy data with PII patterns
+"""
+
+from __future__ import annotations
+
+import hashlib
+import random
+import string
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Any
+
+import polars as pl
+
+
+# =============================================================================
+# Configuration
+# =============================================================================
+
+
+class ColumnType(str, Enum):
+    """Types of columns for data generation."""
+
+    INTEGER = "integer"
+    FLOAT = "float"
+    STRING = "string"
+    BOOLEAN = "boolean"
+    DATE = "date"
+    DATETIME = "datetime"
+    CATEGORY = "category"
+    UUID = "uuid"
+    EMAIL = "email"
+    PHONE = "phone"
+    IP_ADDRESS = "ip_address"
+    URL = "url"
+    JSON = "json"
+
+
+class DataPattern(str, Enum):
+    """Data patterns for generation."""
+
+    UNIFORM = "uniform"
+    NORMAL = "normal"
+    EXPONENTIAL = "exponential"
+    SKEWED = "skewed"
+    SEQUENTIAL = "sequential"
+    CYCLIC = "cyclic"
+    RANDOM_WALK = "random_walk"
+    SPARSE = "sparse"  # Many nulls
+    CONSTANT = "constant"
+
+
+@dataclass
+class ColumnSpec:
+    """Specification for a generated column."""
+
+    name: str
+    dtype: ColumnType
+    pattern: DataPattern = DataPattern.UNIFORM
+    null_ratio: float = 0.0
+    cardinality: int = 0  # 0 = unlimited
+    min_value: float | None = None
+    max_value: float | None = None
+    categories: list[str] | None = None
+    params: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class GeneratorConfig:
+    """Configuration for data generation.
+
+    Controls output size, randomness, and data characteristics.
+    """
+
+    row_count: int = 100_000
+    seed: int | None = None
+    null_ratio: float = 0.01
+    duplicate_ratio: float = 0.0
+    column_specs: list[ColumnSpec] = field(default_factory=list)
+
+    # Performance
+    chunk_size: int = 100_000
+    parallelize: bool = True
+
+    def with_size(self, row_count: int) -> "GeneratorConfig":
+        """Create a copy with different row count."""
+        return GeneratorConfig(
+            row_count=row_count,
+            seed=self.seed,
+            null_ratio=self.null_ratio,
+            duplicate_ratio=self.duplicate_ratio,
+            column_specs=self.column_specs,
+            chunk_size=self.chunk_size,
+            parallelize=self.parallelize,
+        )
+
+
+# =============================================================================
+# Base Generator
+# =============================================================================
+
+
+class DataGenerator(ABC):
+    """Abstract base class for data generators.
+
+    Provides the framework for creating benchmark data.
+
+    Example:
+        class MyGenerator(DataGenerator):
+            name = "my_generator"
+            description = "Generates my custom data"
+
+            def generate(self, config):
+                return pl.DataFrame({...})
+    """
+
+    name: str = "base"
+    description: str = ""
+
+    def __init__(self, seed: int | None = None):
+        self._seed = seed
+        self._rng = random.Random(seed)
+
+    @abstractmethod
+    def generate(self, config: GeneratorConfig) -> pl.DataFrame:
+        """Generate a DataFrame according to the configuration.
+
+        Args:
+            config: Generation configuration
+
+        Returns:
+            Generated DataFrame
+        """
+        pass
+
+    def generate_lazy(self, config: GeneratorConfig) -> pl.LazyFrame:
+        """Generate a LazyFrame.
+
+        Default implementation wraps generate(). Override for
+        streaming generation.
+
+        Args:
+            config: Generation configuration
+
+        Returns:
+            Generated LazyFrame
+        """
+        return self.generate(config).lazy()
+
+    def _apply_nulls(
+        self,
+        series: pl.Series,
+        null_ratio: float,
+    ) -> pl.Series:
+        """Apply null values to a series."""
+        if null_ratio <= 0:
+            return series
+
+        n = len(series)
+        null_count = int(n * null_ratio)
+
+        if null_count == 0:
+            return series
+
+        # Create mask for null positions
+        null_indices = set(self._rng.sample(range(n), null_count))
+
+        # Convert to list, apply nulls, and create new series
+        values = series.to_list()
+        for idx in null_indices:
+            values[idx] = None
+
+        return pl.Series(series.name or "", values)
+
+    def _generate_string(
+        self,
+        n: int,
+        min_length: int = 5,
+        max_length: int = 20,
+    ) -> list[str]:
+        """Generate random strings."""
+        return [
+            "".join(
+                self._rng.choices(
+                    string.ascii_letters + string.digits,
+                    k=self._rng.randint(min_length, max_length),
+                )
+            )
+            for _ in range(n)
+        ]
+
+    def _generate_emails(self, n: int) -> list[str]:
+        """Generate random email addresses."""
+        domains = ["example.com", "test.org", "sample.net", "demo.io"]
+        return [
+            f"user{i}@{self._rng.choice(domains)}"
+            for i in range(n)
+        ]
+
+    def _generate_phones(self, n: int) -> list[str]:
+        """Generate random phone numbers."""
+        return [
+            f"+1-{self._rng.randint(200, 999)}-{self._rng.randint(200, 999)}-{self._rng.randint(1000, 9999)}"
+            for _ in range(n)
+        ]
+
+    def _generate_uuids(self, n: int) -> list[str]:
+        """Generate random UUIDs."""
+        import uuid
+        return [str(uuid.uuid4()) for _ in range(n)]
+
+    def _generate_ips(self, n: int) -> list[str]:
+        """Generate random IP addresses."""
+        return [
+            f"{self._rng.randint(1, 255)}.{self._rng.randint(0, 255)}."
+            f"{self._rng.randint(0, 255)}.{self._rng.randint(1, 254)}"
+            for _ in range(n)
+        ]
+
+
+# =============================================================================
+# Tabular Data Generator
+# =============================================================================
+
+
+class TabularDataGenerator(DataGenerator):
+    """Generator for general-purpose tabular data.
+
+    Creates realistic tabular data with configurable columns,
+    distributions, and data quality characteristics.
+
+    Example:
+        generator = TabularDataGenerator()
+        config = GeneratorConfig(
+            row_count=1_000_000,
+            column_specs=[
+                ColumnSpec("id", ColumnType.INTEGER, DataPattern.SEQUENTIAL),
+                ColumnSpec("value", ColumnType.FLOAT, DataPattern.NORMAL),
+                ColumnSpec("category", ColumnType.CATEGORY, categories=["A", "B", "C"]),
+            ],
+        )
+        df = generator.generate(config)
+    """
+
+    name = "tabular"
+    description = "General-purpose tabular data generator"
+
+    def generate(self, config: GeneratorConfig) -> pl.DataFrame:
+        """Generate tabular data."""
+        if config.seed is not None:
+            self._rng = random.Random(config.seed)
+
+        n = config.row_count
+
+        # Use default schema if no specs provided
+        if not config.column_specs:
+            return self._generate_default_schema(n, config.null_ratio)
+
+        # Generate each column according to spec
+        columns: dict[str, pl.Series] = {}
+        for spec in config.column_specs:
+            series = self._generate_column(n, spec)
+
+            # Apply nulls
+            null_ratio = spec.null_ratio if spec.null_ratio > 0 else config.null_ratio
+            if null_ratio > 0:
+                series = self._apply_nulls(series, null_ratio)
+
+            columns[spec.name] = series
+
+        df = pl.DataFrame(columns)
+
+        # Apply duplicates if configured
+        if config.duplicate_ratio > 0:
+            df = self._add_duplicates(df, config.duplicate_ratio)
+
+        return df
+
+    def _generate_default_schema(
+        self,
+        n: int,
+        null_ratio: float,
+    ) -> pl.DataFrame:
+        """Generate data with a default schema."""
+        return pl.DataFrame({
+            "id": pl.arange(0, n, eager=True),
+            "int_col": pl.Series([self._rng.randint(0, 1000) for _ in range(n)]),
+            "float_col": pl.Series([self._rng.gauss(100, 25) for _ in range(n)]),
+            "string_col": pl.Series(self._generate_string(n)),
+            "category_col": pl.Series(self._rng.choices(["A", "B", "C", "D"], k=n)),
+            "bool_col": pl.Series(self._rng.choices([True, False], k=n)),
+            "date_col": pl.datetime_range(
+                datetime(2020, 1, 1),
+                datetime(2024, 12, 31),
+                eager=True,
+            ).sample(n, with_replacement=True, seed=self._seed),
+        })
+
+    def _generate_column(self, n: int, spec: ColumnSpec) -> pl.Series:
+        """Generate a single column according to spec."""
+        if spec.dtype == ColumnType.INTEGER:
+            return self._generate_integer(n, spec)
+        elif spec.dtype == ColumnType.FLOAT:
+            return self._generate_float(n, spec)
+        elif spec.dtype == ColumnType.STRING:
+            return pl.Series(self._generate_string(n))
+        elif spec.dtype == ColumnType.BOOLEAN:
+            return pl.Series(self._rng.choices([True, False], k=n))
+        elif spec.dtype == ColumnType.CATEGORY:
+            categories = spec.categories or ["cat_A", "cat_B", "cat_C"]
+            return pl.Series(self._rng.choices(categories, k=n))
+        elif spec.dtype == ColumnType.EMAIL:
+            return pl.Series(self._generate_emails(n))
+        elif spec.dtype == ColumnType.PHONE:
+            return pl.Series(self._generate_phones(n))
+        elif spec.dtype == ColumnType.UUID:
+            return pl.Series(self._generate_uuids(n))
+        elif spec.dtype == ColumnType.IP_ADDRESS:
+            return pl.Series(self._generate_ips(n))
+        elif spec.dtype == ColumnType.DATE:
+            return self._generate_dates(n, spec)
+        elif spec.dtype == ColumnType.DATETIME:
+            return self._generate_datetimes(n, spec)
+        else:
+            return pl.Series(self._generate_string(n))
+
+    def _generate_integer(self, n: int, spec: ColumnSpec) -> pl.Series:
+        """Generate integer values according to pattern."""
+        min_val = int(spec.min_value or 0)
+        max_val = int(spec.max_value or 1_000_000)
+
+        if spec.pattern == DataPattern.SEQUENTIAL:
+            return pl.Series(list(range(min_val, min_val + n)))
+        elif spec.pattern == DataPattern.NORMAL:
+            mean = (min_val + max_val) / 2
+            std = (max_val - min_val) / 6
+            values = [int(max(min_val, min(max_val, self._rng.gauss(mean, std)))) for _ in range(n)]
+            return pl.Series(values)
+        elif spec.pattern == DataPattern.EXPONENTIAL:
+            values = [int(min_val + self._rng.expovariate(1 / (max_val - min_val)) % (max_val - min_val)) for _ in range(n)]
+            return pl.Series(values)
+        else:  # UNIFORM
+            return pl.Series([self._rng.randint(min_val, max_val) for _ in range(n)])
+
+    def _generate_float(self, n: int, spec: ColumnSpec) -> pl.Series:
+        """Generate float values according to pattern."""
+        min_val = spec.min_value or 0.0
+        max_val = spec.max_value or 1000.0
+
+        if spec.pattern == DataPattern.NORMAL:
+            mean = (min_val + max_val) / 2
+            std = (max_val - min_val) / 6
+            values = [max(min_val, min(max_val, self._rng.gauss(mean, std))) for _ in range(n)]
+            return pl.Series(values)
+        elif spec.pattern == DataPattern.EXPONENTIAL:
+            scale = (max_val - min_val) / 3
+            values = [min_val + self._rng.expovariate(1 / scale) for _ in range(n)]
+            return pl.Series(values)
+        elif spec.pattern == DataPattern.RANDOM_WALK:
+            values = [min_val + (max_val - min_val) / 2]
+            for _ in range(n - 1):
+                step = self._rng.gauss(0, (max_val - min_val) / 100)
+                values.append(max(min_val, min(max_val, values[-1] + step)))
+            return pl.Series(values)
+        else:  # UNIFORM
+            return pl.Series([self._rng.uniform(min_val, max_val) for _ in range(n)])
+
+    def _generate_dates(self, n: int, spec: ColumnSpec) -> pl.Series:
+        """Generate date values."""
+        start = datetime(2020, 1, 1)
+        end = datetime(2024, 12, 31)
+        delta = (end - start).days
+
+        if spec.pattern == DataPattern.SEQUENTIAL:
+            dates = [start + timedelta(days=i % delta) for i in range(n)]
+        else:
+            dates = [start + timedelta(days=self._rng.randint(0, delta)) for _ in range(n)]
+
+        return pl.Series(dates).cast(pl.Date)
+
+    def _generate_datetimes(self, n: int, spec: ColumnSpec) -> pl.Series:
+        """Generate datetime values."""
+        start = datetime(2020, 1, 1)
+        end = datetime(2024, 12, 31)
+        delta_seconds = int((end - start).total_seconds())
+
+        if spec.pattern == DataPattern.SEQUENTIAL:
+            interval = delta_seconds // n
+            datetimes = [start + timedelta(seconds=i * interval) for i in range(n)]
+        else:
+            datetimes = [start + timedelta(seconds=self._rng.randint(0, delta_seconds)) for _ in range(n)]
+
+        return pl.Series(datetimes)
+
+    def _add_duplicates(
+        self,
+        df: pl.DataFrame,
+        ratio: float,
+    ) -> pl.DataFrame:
+        """Add duplicate rows to the DataFrame."""
+        n = len(df)
+        dup_count = int(n * ratio)
+
+        if dup_count == 0:
+            return df
+
+        # Sample rows to duplicate
+        dup_indices = self._rng.sample(range(n), min(dup_count, n))
+        duplicates = df[dup_indices]
+
+        return pl.concat([df, duplicates])
+
+
+# =============================================================================
+# Time Series Data Generator
+# =============================================================================
+
+
+class TimeSeriesDataGenerator(DataGenerator):
+    """Generator for time-series data with patterns.
+
+    Creates time-series data with trends, seasonality, and noise.
+
+    Example:
+        generator = TimeSeriesDataGenerator()
+        config = GeneratorConfig(row_count=100_000)
+        df = generator.generate(config)
+    """
+
+    name = "timeseries"
+    description = "Time-series data with patterns"
+
+    def __init__(
+        self,
+        seed: int | None = None,
+        frequency: str = "1h",
+        trend: float = 0.0,
+        seasonality_period: int = 24,
+        seasonality_amplitude: float = 10.0,
+        noise_std: float = 5.0,
+    ):
+        super().__init__(seed)
+        self.frequency = frequency
+        self.trend = trend
+        self.seasonality_period = seasonality_period
+        self.seasonality_amplitude = seasonality_amplitude
+        self.noise_std = noise_std
+
+    def generate(self, config: GeneratorConfig) -> pl.DataFrame:
+        """Generate time-series data."""
+        if config.seed is not None:
+            self._rng = random.Random(config.seed)
+
+        n = config.row_count
+        import math
+
+        # Generate timestamps
+        start = datetime(2024, 1, 1)
+        timestamps = pl.datetime_range(
+            start,
+            start + timedelta(hours=n),
+            interval=self.frequency,
+            eager=True,
+        ).head(n)
+
+        # Generate values with trend + seasonality + noise
+        values = []
+        for i in range(n):
+            trend_component = self.trend * i
+            seasonal_component = self.seasonality_amplitude * math.sin(
+                2 * math.pi * i / self.seasonality_period
+            )
+            noise_component = self._rng.gauss(0, self.noise_std)
+            value = 100 + trend_component + seasonal_component + noise_component
+            values.append(value)
+
+        return pl.DataFrame({
+            "timestamp": timestamps,
+            "value": values,
+            "hour": [(i % 24) for i in range(n)],
+            "day_of_week": [(i // 24) % 7 for i in range(n)],
+        })
+
+
+# =============================================================================
+# Financial Data Generator
+# =============================================================================
+
+
+class FinancialDataGenerator(DataGenerator):
+    """Generator for financial market data.
+
+    Creates realistic OHLCV (Open, High, Low, Close, Volume) data
+    for stocks and cryptocurrencies.
+
+    Example:
+        generator = FinancialDataGenerator(symbol="AAPL")
+        config = GeneratorConfig(row_count=1_000_000)
+        df = generator.generate(config)
+    """
+
+    name = "financial"
+    description = "Stock/crypto OHLCV data"
+
+    def __init__(
+        self,
+        seed: int | None = None,
+        symbol: str = "AAPL",
+        base_price: float = 100.0,
+        volatility: float = 0.02,
+        frequency: str = "1s",
+    ):
+        super().__init__(seed)
+        self.symbol = symbol
+        self.base_price = base_price
+        self.volatility = volatility
+        self.frequency = frequency
+
+    def generate(self, config: GeneratorConfig) -> pl.DataFrame:
+        """Generate OHLCV data."""
+        if config.seed is not None:
+            self._rng = random.Random(config.seed)
+
+        n = config.row_count
+
+        # Generate timestamps
+        start = datetime(2024, 1, 1)
+        timestamps = pl.datetime_range(
+            start,
+            start + timedelta(seconds=n),
+            interval=self.frequency,
+            eager=True,
+        ).head(n)
+
+        # Generate price using geometric Brownian motion
+        prices = [self.base_price]
+        for _ in range(n - 1):
+            drift = 0.0001  # Small upward drift
+            shock = self._rng.gauss(0, self.volatility)
+            prices.append(prices[-1] * (1 + drift + shock))
+
+        # Generate OHLCV
+        opens = prices
+        highs = [p * (1 + abs(self._rng.gauss(0, self.volatility / 2))) for p in prices]
+        lows = [p * (1 - abs(self._rng.gauss(0, self.volatility / 2))) for p in prices]
+        closes = [
+            lows[i] + self._rng.random() * (highs[i] - lows[i])
+            for i in range(n)
+        ]
+        volumes = [int(self._rng.expovariate(1 / 100000)) for _ in range(n)]
+
+        return pl.DataFrame({
+            "timestamp": timestamps,
+            "symbol": [self.symbol] * n,
+            "open": opens,
+            "high": highs,
+            "low": lows,
+            "close": closes,
+            "volume": volumes,
+            "vwap": [(h + l + c) / 3 for h, l, c in zip(highs, lows, closes)],
+        })
+
+
+# =============================================================================
+# Text Data Generator
+# =============================================================================
+
+
+class TextDataGenerator(DataGenerator):
+    """Generator for text-heavy data with PII patterns.
+
+    Creates data with emails, names, addresses, and other
+    PII-like patterns for testing scanners.
+
+    Example:
+        generator = TextDataGenerator()
+        config = GeneratorConfig(row_count=100_000)
+        df = generator.generate(config)
+    """
+
+    name = "text"
+    description = "Text data with PII patterns"
+
+    # Sample data pools
+    FIRST_NAMES = [
+        "James", "Mary", "John", "Patricia", "Robert", "Jennifer",
+        "Michael", "Linda", "William", "Elizabeth", "David", "Barbara",
+    ]
+    LAST_NAMES = [
+        "Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia",
+        "Miller", "Davis", "Rodriguez", "Martinez", "Hernandez", "Lopez",
+    ]
+    DOMAINS = ["gmail.com", "yahoo.com", "outlook.com", "company.com"]
+    CITIES = ["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"]
+    STATES = ["NY", "CA", "IL", "TX", "AZ"]
+
+    def generate(self, config: GeneratorConfig) -> pl.DataFrame:
+        """Generate text data with PII."""
+        if config.seed is not None:
+            self._rng = random.Random(config.seed)
+
+        n = config.row_count
+
+        first_names = [self._rng.choice(self.FIRST_NAMES) for _ in range(n)]
+        last_names = [self._rng.choice(self.LAST_NAMES) for _ in range(n)]
+
+        return pl.DataFrame({
+            "id": pl.arange(0, n, eager=True),
+            "first_name": first_names,
+            "last_name": last_names,
+            "email": [
+                f"{fn.lower()}.{ln.lower()}{self._rng.randint(1, 999)}@{self._rng.choice(self.DOMAINS)}"
+                for fn, ln in zip(first_names, last_names)
+            ],
+            "phone": self._generate_phones(n),
+            "ssn": [
+                f"{self._rng.randint(100, 999)}-{self._rng.randint(10, 99)}-{self._rng.randint(1000, 9999)}"
+                for _ in range(n)
+            ],
+            "credit_card": [
+                f"{self._rng.randint(4000, 5999)}-{self._rng.randint(1000, 9999)}-"
+                f"{self._rng.randint(1000, 9999)}-{self._rng.randint(1000, 9999)}"
+                for _ in range(n)
+            ],
+            "ip_address": self._generate_ips(n),
+            "address": [
+                f"{self._rng.randint(100, 9999)} {self._rng.choice(['Main', 'Oak', 'Pine', 'Elm'])} "
+                f"{self._rng.choice(['St', 'Ave', 'Blvd', 'Rd'])}"
+                for _ in range(n)
+            ],
+            "city": [self._rng.choice(self.CITIES) for _ in range(n)],
+            "state": [self._rng.choice(self.STATES) for _ in range(n)],
+            "zip_code": [f"{self._rng.randint(10000, 99999)}" for _ in range(n)],
+        })
+
+
+# =============================================================================
+# Registry
+# =============================================================================
+
+
+class GeneratorRegistry:
+    """Registry for data generator discovery and management."""
+
+    def __init__(self) -> None:
+        self._generators: dict[str, type[DataGenerator]] = {}
+        self._lock = threading.Lock()
+
+    def register(
+        self,
+        generator_class: type[DataGenerator],
+    ) -> type[DataGenerator]:
+        """Register a generator class."""
+        with self._lock:
+            name = generator_class.name
+            if name in self._generators:
+                raise ValueError(f"Generator '{name}' is already registered")
+            self._generators[name] = generator_class
+        return generator_class
+
+    def get(self, name: str) -> type[DataGenerator]:
+        """Get a generator class by name."""
+        with self._lock:
+            if name not in self._generators:
+                available = list(self._generators.keys())
+                raise KeyError(
+                    f"Generator '{name}' not found. Available: {available}"
+                )
+            return self._generators[name]
+
+    def list_names(self) -> list[str]:
+        """List all registered generator names."""
+        with self._lock:
+            return list(self._generators.keys())
+
+    def create(self, name: str, **kwargs: Any) -> DataGenerator:
+        """Create a generator instance by name."""
+        generator_class = self.get(name)
+        return generator_class(**kwargs)
+
+
+# Global registry instance
+generator_registry = GeneratorRegistry()
+
+# Register built-in generators
+generator_registry.register(TabularDataGenerator)
+generator_registry.register(TimeSeriesDataGenerator)
+generator_registry.register(FinancialDataGenerator)
+generator_registry.register(TextDataGenerator)
+
+
+def register_generator(cls: type[DataGenerator]) -> type[DataGenerator]:
+    """Decorator to register a generator with the global registry."""
+    return generator_registry.register(cls)