benchbox 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

benchbox/platforms/base/cloud_spark/session.py

@@ -41,6 +41,8 @@ from dataclasses import dataclass, field
 from enum import Enum
 from typing import TYPE_CHECKING, Any
 
+from benchbox.utils.dependencies import get_package_install_message
+
 if TYPE_CHECKING:
     from pyspark.sql import SparkSession
 
@@ -332,7 +334,7 @@ class LivySessionManager(CloudSparkSessionManager):
             try:
                 import requests
             except ImportError as e:
-                raise ImportError("requests required for Livy API. Install with: uv add requests") from e
+                raise ImportError(get_package_install_message("requests", "requests required for Livy API.")) from e
             self._http_client = requests.Session()
         return self._http_client
 
@@ -495,7 +497,7 @@ class DatabricksConnectSessionManager(CloudSparkSessionManager):
         try:
             from databricks.connect import DatabricksSession
         except ImportError as e:
-            raise ImportError("databricks-connect required. Install with: uv add databricks-connect") from e
+            raise ImportError(get_package_install_message("databricks-connect", "databricks-connect required.")) from e
 
         self._state = SessionState.STARTING
         self._logger.info("Creating Databricks Connect session...")

benchbox/platforms/base/cloud_spark/staging.py

@@ -40,6 +40,8 @@ from pathlib import Path
 from typing import Any, Callable
 from urllib.parse import urlparse
 
+from benchbox.utils.dependencies import get_package_install_message
+
 logger = logging.getLogger(__name__)
 
 
@@ -370,7 +372,7 @@ class S3Staging(CloudSparkStaging):
             try:
                 import boto3
             except ImportError as e:
-                raise ImportError("boto3 required for S3 staging. Install with: uv add boto3") from e
+                raise ImportError(get_package_install_message("boto3", "boto3 required for S3 staging.")) from e
 
             session_kwargs = {}
             if self.config.credentials:
@@ -462,7 +464,9 @@ class GCSStaging(CloudSparkStaging):
                 from google.cloud import storage
             except ImportError as e:
                 raise ImportError(
-                    "google-cloud-storage required for GCS staging. Install with: uv add google-cloud-storage"
+                    get_package_install_message(
+                        "google-cloud-storage", "google-cloud-storage required for GCS staging."
+                    )
                 ) from e
 
             self._client = storage.Client()
@@ -539,8 +543,10 @@ class AzureADLSStaging(CloudSparkStaging):
                 from azure.storage.filedatalake import DataLakeServiceClient
             except ImportError as e:
                 raise ImportError(
-                    "azure-storage-file-datalake required for Azure ADLS staging. "
-                    "Install with: uv add azure-storage-file-datalake azure-identity"
+                    get_package_install_message(
+                        "azure-storage-file-datalake azure-identity",
+                        "azure-storage-file-datalake required for Azure ADLS staging.",
+                    )
                 ) from e
 
             # Parse account from container@account.dfs.core.windows.net
@@ -628,8 +634,10 @@ class AzureBlobStaging(CloudSparkStaging):
                 from azure.storage.blob import ContainerClient
             except ImportError as e:
                 raise ImportError(
-                    "azure-storage-blob required for Azure Blob staging. "
-                    "Install with: uv add azure-storage-blob azure-identity"
+                    get_package_install_message(
+                        "azure-storage-blob azure-identity",
+                        "azure-storage-blob required for Azure Blob staging.",
+                    )
                 ) from e
 
             credential = DefaultAzureCredential()
@@ -707,7 +715,7 @@ class DBFSStaging(CloudSparkStaging):
                 from databricks.sdk import WorkspaceClient
             except ImportError as e:
                 raise ImportError(
-                    "databricks-sdk required for DBFS staging. Install with: uv add databricks-sdk"
+                    get_package_install_message("databricks-sdk", "databricks-sdk required for DBFS staging.")
                 ) from e
 
             self._client = WorkspaceClient()

benchbox/platforms/base/data_loading.py

@@ -22,6 +22,8 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Protocol, cast
 
+from benchbox.utils.file_format import detect_compression
+
 logger = logging.getLogger(__name__)
 
 # Regex pattern for valid SQL identifiers (table/column names)
@@ -944,6 +946,177 @@ class DuckDBDeltaHandler(FileFormatHandler):
             return row_count
 
 
+class DuckLakeFileHandler(FileFormatHandler):
+    """Handler for DuckLake tables using DuckDB's ducklake extension.
+
+    This handler reads DuckLake tables and loads data via INSERT statements.
+    """
+
+    def get_delimiter(self) -> str:
+        """DuckLake is a table format, not delimited."""
+        return ""  # Not applicable for DuckLake
+
+    def load_table(self, table_name: str, file_path: Path, connection: Any, benchmark: Any, logger: Any) -> int:
+        """Load DuckLake table into database table.
+
+        Args:
+            table_name: Name of table to load
+            file_path: Path to DuckLake table directory
+            connection: Database connection
+            benchmark: Benchmark instance
+            logger: Logger instance
+
+        Returns:
+            Number of rows loaded
+        """
+        # Validate table name to prevent SQL injection
+        validated_table = validate_sql_identifier(table_name, "table name")
+
+        try:
+            import duckdb
+        except ImportError as e:
+            raise RuntimeError("DuckLake support requires DuckDB. Install it with: uv add duckdb") from e
+
+        # DuckLake tables have a metadata.ducklake file and data directory
+        metadata_path = file_path / "metadata.ducklake"
+        data_path = file_path / "data"
+
+        if not metadata_path.exists():
+            raise RuntimeError(f"DuckLake catalog not found at {metadata_path}")
+
+        # Create a temporary connection to read the DuckLake table
+        temp_conn = duckdb.connect(":memory:")
+        try:
+            # Load ducklake extension
+            temp_conn.execute("INSTALL ducklake")
+            temp_conn.execute("LOAD ducklake")
+
+            # Attach the DuckLake catalog (use DATA_PATH option, not query string)
+            temp_conn.execute(f"ATTACH 'ducklake:{metadata_path}' AS ducklake_db (DATA_PATH '{data_path}')")
+
+            # Read data from DuckLake table
+            arrow_table = temp_conn.execute(f"SELECT * FROM ducklake_db.main.{validated_table}").fetch_arrow_table()
+
+            row_count = arrow_table.num_rows
+
+            if row_count == 0:
+                return 0
+
+            # Convert to Python data for insertion
+            data = arrow_table.to_pylist()
+
+            # Get column names from schema and validate each one
+            column_names = arrow_table.schema.names
+            validated_columns = [validate_sql_identifier(col, "column name") for col in column_names]
+            placeholders = ",".join(["?" for _ in validated_columns])
+            columns_str = ",".join(validated_columns)
+
+            # Prepare INSERT statement
+            insert_sql = f"INSERT INTO {validated_table} ({columns_str}) VALUES ({placeholders})"
+
+            # Convert dicts to tuples in correct column order
+            data_tuples = [tuple(row[col] for col in column_names) for row in data]
+
+            # Insert in batches for better performance
+            batch_size = 1000
+            for i in range(0, len(data_tuples), batch_size):
+                batch = data_tuples[i : i + batch_size]
+                connection.executemany(insert_sql, batch)
+
+            return row_count
+
+        finally:
+            temp_conn.close()
+
+
+class DuckDBDuckLakeHandler(FileFormatHandler):
+    """Handler for DuckDB's native DuckLake support.
+
+    This handler leverages DuckDB's ducklake extension for optimized
+    DuckLake reading directly within DuckDB.
+    """
+
+    def __init__(self, adapter: Any):
+        """Initialize handler.
+
+        Args:
+            adapter: Platform adapter (for dry-run support)
+        """
+        self.adapter = adapter
+
+    def get_delimiter(self) -> str:
+        """DuckLake is a table format, not delimited."""
+        return ""
+
+    def load_table(self, table_name: str, file_path: Path, connection: Any, benchmark: Any, logger: Any) -> int:
+        """Load table using DuckDB's ducklake extension.
+
+        Args:
+            table_name: Name of table to load
+            file_path: Path to DuckLake table directory
+            connection: DuckDB connection
+            benchmark: Benchmark instance
+            logger: Logger instance
+
+        Returns:
+            Number of rows loaded
+        """
+        # Validate table name to prevent SQL injection
+        validated_table = validate_sql_identifier(table_name, "table name")
+
+        # DuckLake tables have a metadata.ducklake file and data directory
+        metadata_path = file_path / "metadata.ducklake"
+        data_path = file_path / "data"
+
+        if not metadata_path.exists():
+            raise RuntimeError(f"DuckLake catalog not found at {metadata_path}")
+
+        # Ensure ducklake extension is installed and loaded
+        try:
+            connection.execute("INSTALL ducklake")
+            connection.execute("LOAD ducklake")
+        except Exception:
+            # Extension might already be installed/loaded
+            pass
+
+        # Generate a unique alias for the catalog
+        import uuid
+
+        catalog_alias = f"ducklake_{uuid.uuid4().hex[:8]}"
+
+        # Attach the DuckLake catalog (use DATA_PATH option, not query string)
+        escaped_metadata = escape_sql_string_literal(str(metadata_path))
+        escaped_data_path = escape_sql_string_literal(str(data_path))
+
+        try:
+            connection.execute(
+                f"ATTACH 'ducklake:{escaped_metadata}' AS {catalog_alias} (DATA_PATH '{escaped_data_path}')"
+            )
+
+            # Build INSERT SELECT from the DuckLake table
+            insert_sql = f"""
+                INSERT INTO {validated_table}
+                SELECT * FROM {catalog_alias}.main.{validated_table}
+            """
+
+            # Execute or capture based on dry-run mode
+            if hasattr(self.adapter, "dry_run_mode") and self.adapter.dry_run_mode:
+                self.adapter.capture_sql(insert_sql, "load_data", validated_table)
+                return 1000  # Placeholder for dry-run
+            else:
+                connection.execute(insert_sql)
+                # Get actual row count
+                row_count = connection.execute(f"SELECT COUNT(*) FROM {validated_table}").fetchone()[0]
+                return row_count
+
+        finally:
+            # Detach the catalog
+            try:
+                connection.execute(f"DETACH {catalog_alias}")
+            except Exception:
+                pass
+
+
 class IcebergFileHandler(FileFormatHandler):
     """Handler for Apache Iceberg tables using pyiceberg library.
 
@@ -1052,6 +1225,137 @@ class IcebergFileHandler(FileFormatHandler):
                     pass
 
 
+class VortexFileHandler(FileFormatHandler):
+    """Handler for Vortex columnar files using the vortex Python library.
+
+    This is a generic handler that works across platforms by reading
+    Vortex files and loading via INSERT statements.
+    """
+
+    def get_delimiter(self) -> str:
+        """Vortex is a columnar format, not delimited."""
+        return ""  # Not applicable for Vortex
+
+    def load_table(self, table_name: str, file_path: Path, connection: Any, benchmark: Any, logger: Any) -> int:
+        """Load Vortex file into table.
+
+        Args:
+            table_name: Name of table to load
+            file_path: Path to Vortex file
+            connection: Database connection
+            benchmark: Benchmark instance
+            logger: Logger instance
+
+        Returns:
+            Number of rows loaded
+        """
+        # Validate table name to prevent SQL injection
+        validated_table = validate_sql_identifier(table_name, "table name")
+
+        try:
+            import vortex
+        except ImportError as e:
+            raise RuntimeError(
+                "Vortex format support requires the 'vortex' package. "
+                "Install it with: uv add vortex --optional table-formats"
+            ) from e
+
+        # Read Vortex file and convert to PyArrow table
+        vortex_array = vortex.io.read(file_path)
+        arrow_table = vortex_array.to_arrow()
+        row_count = arrow_table.num_rows
+
+        if row_count == 0:
+            return 0
+
+        # Convert to Python data for insertion
+        data = arrow_table.to_pylist()
+
+        # Get column names from schema and validate each one
+        column_names = arrow_table.schema.names
+        validated_columns = [validate_sql_identifier(col, "column name") for col in column_names]
+        placeholders = ",".join(["?" for _ in validated_columns])
+        columns_str = ",".join(validated_columns)
+
+        # Prepare INSERT statement
+        insert_sql = f"INSERT INTO {validated_table} ({columns_str}) VALUES ({placeholders})"
+
+        # Convert dicts to tuples in correct column order
+        data_tuples = [tuple(row[col] for col in column_names) for row in data]
+
+        # Insert in batches for better performance
+        batch_size = 1000
+        for i in range(0, len(data_tuples), batch_size):
+            batch = data_tuples[i : i + batch_size]
+            connection.executemany(insert_sql, batch)
+
+        return row_count
+
+
+class DuckDBVortexHandler(FileFormatHandler):
+    """Handler for DuckDB's native Vortex support via extension.
+
+    This handler leverages DuckDB's vortex extension for optimized
+    Vortex reading instead of using the Python vortex library.
+    """
+
+    def __init__(self, adapter: Any):
+        """Initialize handler.
+
+        Args:
+            adapter: Platform adapter (for dry-run support)
+        """
+        self.adapter = adapter
+
+    def get_delimiter(self) -> str:
+        """Vortex is a columnar format, not delimited."""
+        return ""
+
+    def load_table(self, table_name: str, file_path: Path, connection: Any, benchmark: Any, logger: Any) -> int:
+        """Load table using DuckDB's vortex extension.
+
+        Args:
+            table_name: Name of table to load
+            file_path: Path to Vortex file
+            connection: DuckDB connection
+            benchmark: Benchmark instance
+            logger: Logger instance
+
+        Returns:
+            Number of rows loaded
+        """
+        # Validate table name and escape file path to prevent SQL injection
+        validated_table = validate_sql_identifier(table_name, "table name")
+        escaped_path = escape_sql_string_literal(str(file_path))
+
+        # Ensure vortex extension is installed and loaded
+        try:
+            # Install vortex extension if not already installed
+            connection.execute("INSTALL vortex")
+            connection.execute("LOAD vortex")
+        except Exception:
+            # Extension might already be installed/loaded, or not available
+            # Fall back to generic handler
+            logger.debug("DuckDB vortex extension not available, falling back to generic handler")
+            return VortexFileHandler().load_table(table_name, file_path, connection, benchmark, logger)
+
+        # Build INSERT SELECT from read_vortex()
+        insert_sql = f"""
+            INSERT INTO {validated_table}
+            SELECT * FROM read_vortex('{escaped_path}')
+        """
+
+        # Execute or capture based on dry-run mode
+        if hasattr(self.adapter, "dry_run_mode") and self.adapter.dry_run_mode:
+            self.adapter.capture_sql(insert_sql, "load_data", validated_table)
+            return 1000  # Placeholder for dry-run
+        else:
+            connection.execute(insert_sql)
+            # Get actual row count
+            row_count = connection.execute(f"SELECT COUNT(*) FROM {validated_table}").fetchone()[0]
+            return row_count
+
+
 class ClickHouseNativeHandler(FileFormatHandler):
     """Handler for ClickHouse's native file() function.
 
@@ -1137,7 +1441,7 @@ class ClickHouseNativeHandler(FileFormatHandler):
         actual_file_path = file_path
         temp_file_path = None
 
-        if hasattr(self.adapter, "mode") and self.adapter.mode == "local" and str(file_path).endswith(".zst"):
+        if hasattr(self.adapter, "mode") and self.adapter.mode == "local" and detect_compression(file_path) == "zstd":
             # Decompress zstd files for local mode
             self.adapter.log_verbose(f"Decompressing {file_path.name} for local mode...")
             try:
@@ -1256,6 +1560,7 @@ class FileFormatRegistry:
         ".tbl": lambda: DelimitedFileHandler("|"),
         ".dat": lambda: DelimitedFileHandler("|"),
         ".parquet": lambda: ParquetFileHandler(),
+        ".vortex": lambda: VortexFileHandler(),
     }
 
     # Map compression extensions to handlers
@@ -1306,6 +1611,11 @@ class FileFormatRegistry:
         """
         # Check if this is a directory-based table format
         if file_path.is_dir():
+            # Check for DuckLake (metadata.ducklake file)
+            ducklake_metadata = file_path / "metadata.ducklake"
+            if ducklake_metadata.exists():
+                return DuckLakeFileHandler()
+
             # Check for Delta Lake (_delta_log directory)
             delta_log_dir = file_path / "_delta_log"
             if delta_log_dir.exists() and delta_log_dir.is_dir():
@@ -1540,7 +1850,11 @@ __all__ = [
     "DuckDBParquetHandler",
     "DeltaFileHandler",
     "DuckDBDeltaHandler",
+    "DuckLakeFileHandler",
+    "DuckDBDuckLakeHandler",
     "IcebergFileHandler",
+    "VortexFileHandler",
+    "DuckDBVortexHandler",
     "ClickHouseNativeHandler",
     "InMemoryDataHandler",
     "FileFormatRegistry",

benchbox/platforms/base/format_capabilities.py

@@ -99,17 +99,52 @@ ICEBERG_CAPABILITY = FormatCapability(
     },
 )
 
+VORTEX_CAPABILITY = FormatCapability(
+    format_name="vortex",
+    display_name="Vortex",
+    file_extension=".vortex",
+    features={
+        "compression",
+        "column_pruning",
+        "predicate_pushdown",
+        "statistics",
+    },
+    supported_platforms={
+        "duckdb": SupportLevel.EXTENSION,  # vortex extension
+        "datafusion": SupportLevel.EXPERIMENTAL,
+    },
+)
+
+DUCKLAKE_CAPABILITY = FormatCapability(
+    format_name="ducklake",
+    display_name="DuckLake",
+    file_extension="",  # DuckLake is directory-based
+    features={
+        "time_travel",
+        "acid_transactions",
+        "schema_evolution",
+        "snapshot_isolation",
+        "predicate_pushdown",
+        "column_pruning",
+    },
+    supported_platforms={
+        "duckdb": SupportLevel.NATIVE,  # DuckLake is DuckDB's native table format
+    },
+)
+
 # Registry of all format capabilities
 CAPABILITIES_REGISTRY: dict[str, FormatCapability] = {
     "parquet": PARQUET_CAPABILITY,
     "delta": DELTA_CAPABILITY,
     "iceberg": ICEBERG_CAPABILITY,
+    "vortex": VORTEX_CAPABILITY,
+    "ducklake": DUCKLAKE_CAPABILITY,
 }
 
 # Platform format preferences (which format to prefer when multiple available)
 PLATFORM_FORMAT_PREFERENCES: dict[str, list[str]] = {
-    "duckdb": ["parquet", "delta", "tbl", "csv"],
-    "datafusion": ["parquet", "tbl", "csv"],
+    "duckdb": ["parquet", "ducklake", "vortex", "delta", "tbl", "csv"],
+    "datafusion": ["parquet", "vortex", "tbl", "csv"],
     "clickhouse": ["parquet", "tbl", "csv"],
     "databricks": ["delta", "parquet", "tbl", "csv"],
     "snowflake": ["parquet", "tbl", "csv"],

benchbox/platforms/base/utils.py

@@ -7,6 +7,8 @@ import sys
 from pathlib import Path
 from typing import NamedTuple
 
+from benchbox.utils.file_format import is_parquet_format, is_tpc_format
+
 
 class FileFormatInfo(NamedTuple):
     """Information about detected file format."""
@@ -34,17 +36,17 @@ def detect_file_format(file_paths: list[Path] | Path) -> FileFormatInfo:
     if not file_paths:
         return FileFormatInfo(format_type="csv", delimiter=",", has_trailing_delimiter=False)
 
-    first_file = str(file_paths[0])
+    first_file = file_paths[0]
 
-    # Check for Parquet
-    if first_file.endswith(".parquet"):
+    # Check for Parquet (handles compression like .parquet.zst)
+    if is_parquet_format(first_file):
         return FileFormatInfo(format_type="parquet", delimiter="", has_trailing_delimiter=False)
 
     # Check for TPC benchmark format (.tbl, .dat)
     # BenchBox generators produce TPC files WITHOUT trailing delimiters:
     # - TPC-H: dbgen binary doesn't add trailing delimiters
     # - TPC-DS: dsdgen uses -terminate n flag to disable trailing delimiters
-    if ".tbl" in first_file or ".dat" in first_file:
+    if is_tpc_format(first_file):
         return FileFormatInfo(format_type="tpc", delimiter="|", has_trailing_delimiter=False)
 
     # Default to standard CSV

benchbox/platforms/bigquery.py

@@ -26,6 +26,7 @@ if TYPE_CHECKING:
     )
 
 from benchbox.utils.cloud_storage import get_cloud_path_info, is_cloud_path
+from benchbox.utils.file_format import detect_compression, get_delimiter_for_file
 
 from ..utils.dependencies import check_platform_dependencies, get_dependency_error_message
 from .base import PlatformAdapter
@@ -933,7 +934,7 @@ class BigQueryAdapter(PlatformAdapter):
                 if not isinstance(file_paths, list):
                     file_paths = [file_paths]
                 for file_path in file_paths:
-                    if Path(file_path).suffix == ".zst":
+                    if detect_compression(file_path) == "zstd":
                         benchmark_name = getattr(benchmark, "name", "unknown")
                         scale_factor = getattr(benchmark, "scale_factor", "unknown")
                         raise ValueError(
@@ -996,9 +997,8 @@ class BigQueryAdapter(PlatformAdapter):
 
                             # Detect file format to determine delimiter
                             # TPC-H uses .tbl (pipe-delimited), TPC-DS uses .dat (pipe-delimited)
-                            # Use substring check to handle chunked files like customer.tbl.1 or customer.tbl.1.zst
-                            file_str = str(file_path.name)
-                            delimiter = "|" if ".tbl" in file_str or ".dat" in file_str else ","
+                            # get_delimiter_for_file handles chunked files like customer.tbl.1 or customer.tbl.1.zst
+                            delimiter = get_delimiter_for_file(file_path)
 
                             # Upload file directly to Cloud Storage with original compression and format
                             # BigQuery supports compressed files and any delimiter natively
@@ -1083,8 +1083,7 @@ class BigQueryAdapter(PlatformAdapter):
                             self.log_very_verbose(f"Loading {table_name}{chunk_info} from {file_path.name}")
 
                             # Detect delimiter from filename (handle chunked files like customer.tbl.1)
-                            file_str = str(file_path.name)
-                            delimiter = "|" if ".tbl" in file_str or ".dat" in file_str else ","
+                            delimiter = get_delimiter_for_file(file_path)
 
                             # Use WRITE_APPEND for subsequent chunks, WRITE_TRUNCATE for first
                             write_disposition = (