Source code for benchbox.joinorder

"""Join Order Benchmark top-level interface.

Copyright 2026 Joe Harris / BenchBox Project

Licensed under the MIT License. See LICENSE file in the project root for details.
"""

from pathlib import Path
from typing import Any, Optional, Union

from benchbox.base import BaseBenchmark
from benchbox.core.joinorder.benchmark import JoinOrderBenchmark




class JoinOrder(BaseBenchmark):
    """Join Order Benchmark implementation.

    This class provides an implementation of the Join Order Benchmark, including
    data generation and access to complex join queries for cardinality estimation
    and join order optimization testing.

    Reference: Viktor Leis et al. "How Good Are Query Optimizers, Really?"
    """



    def __init__(
        self,
        scale_factor: float = 1.0,
        output_dir: Optional[Union[str, Path]] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize a Join Order Benchmark instance.

        Args:
            scale_factor: Scale factor for the benchmark (1.0 = ~1GB)
            output_dir: Directory to output generated data files
            **kwargs: Additional implementation-specific options
        """
        super().__init__(scale_factor=scale_factor, output_dir=output_dir, **kwargs)

        # Initialize the actual implementation
        verbose = kwargs.pop("verbose", False)
        self._impl = JoinOrderBenchmark(scale_factor=scale_factor, output_dir=output_dir, verbose=verbose, **kwargs)




    def generate_data(self) -> list[Path]:
        """Generate Join Order Benchmark data.

        Returns:
            A list of paths to the generated data files
        """
        return self._impl.generate_data()




    def get_queries(self) -> dict[str, str]:
        """Get all Join Order Benchmark queries.

        Returns:
            A dictionary mapping query IDs to query strings
        """
        return self._impl.get_queries()




    def get_query(self, query_id: Union[int, str], *, params: Optional[dict[str, Any]] = None) -> str:
        """Get a specific Join Order Benchmark query.

        Args:
            query_id: The ID of the query to retrieve
            params: Optional parameters to customize the query

        Returns:
            The query string

        Raises:
            ValueError: If the query_id is invalid
        """
        return self._impl.get_query(query_id, params=params)




    def get_schema(self, dialect: str = "sqlite") -> str:
        """Get the Join Order Benchmark schema DDL.

        Args:
            dialect: Target SQL dialect

        Returns:
            DDL statements for creating all tables
        """
        return self._impl.get_schema(dialect)




    def get_create_tables_sql(self, dialect: str = "standard", tuning_config: Any = None) -> str:
        """Get SQL to create all Join Order Benchmark tables.

        Args:
            dialect: SQL dialect to use
            tuning_config: Unified tuning configuration for constraint settings

        Returns:
            SQL script for creating all tables
        """
        return self._impl.get_create_tables_sql(dialect=dialect)




__all__ = ["JoinOrder"]