sqlbenchdag 0.1.0__py3-none-any.whl

sql_benchmarks/resources/actian_client.py

@@ -0,0 +1,231 @@
+import os
+import time
+import tempfile
+import polars as pl
+from typing import Dict, Any, Optional
+
+
+class ActianClient:
+    """
+    Stateful Client for Actian Vector on EC2.
+
+    Executes queries via SSH and handles bulk loading via SCP + vwload.
+    """
+
+    def __init__(self, connection_params: Dict[str, Any], ssh_params: Dict[str, Any]):
+        self.connection_params = connection_params
+        self.ssh_params = ssh_params
+        self._ssh_client = None
+
+    def _ensure_ssh(self):
+        """Establishes SSH client for remote commands."""
+        if self._ssh_client is not None:
+            try:
+                transport = self._ssh_client.get_transport()
+                if transport and transport.is_active():
+                    return
+            except Exception:
+                pass
+            self._ssh_client = None
+
+        import paramiko
+
+        self._ssh_client = paramiko.SSHClient()
+        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        self._ssh_client.connect(
+            hostname=self.ssh_params["ec2_host"],
+            username=self.ssh_params["ec2_user"],
+            key_filename=self.ssh_params["ssh_key_path"],
+            timeout=30
+        )
+
+    def _ssh_exec(self, command: str, timeout: int = 300) -> tuple[str, str, int]:
+        """Execute a command on the remote EC2 instance."""
+        self._ensure_ssh()
+
+        stdin, stdout, stderr = self._ssh_client.exec_command(command, timeout=timeout)
+        exit_code = stdout.channel.recv_exit_status()
+
+        return stdout.read().decode(), stderr.read().decode(), exit_code
+
+    def _scp_to_remote(self, local_path: str, remote_path: str):
+        """Copy a file to the remote EC2 instance."""
+        self._ensure_ssh()
+
+        from scp import SCPClient
+
+        # Ensure remote directory exists
+        remote_dir = os.path.dirname(remote_path)
+        self._ssh_exec(f"mkdir -p {remote_dir}")
+
+        with SCPClient(self._ssh_client.get_transport()) as scp:
+            scp.put(local_path, remote_path)
+
+        print(f"[Actian] Uploaded {local_path} -> {remote_path}")
+
+    def run_query(self, sql: str, scenario_params: Dict[str, Any]) -> Optional[float]:
+        """
+        Execute a SQL query on Actian Vector via SSH.
+
+        Uses the native `sql` command on EC2 for reliable execution
+        without requiring local ODBC drivers.
+        """
+        database = self.connection_params["database"]
+        sql_path = self.ssh_params["actian_sql_path"]
+
+        # Escape single quotes in SQL for shell
+        escaped_sql = sql.replace("'", "'\\''")
+
+        # Build command: echo SQL | sql database
+        # The -s flag suppresses headers for cleaner output
+        command = f"echo '{escaped_sql}' | {sql_path} {database}"
+
+        start = time.time()
+        stdout, stderr, exit_code = self._ssh_exec(command, timeout=600)
+        duration = time.time() - start
+
+        if exit_code != 0:
+            raise RuntimeError(f"Actian query failed: {stderr}\nSQL: {sql[:200]}...")
+
+        return duration
+
+    def bulk_load(self, file_path: str, table_name: str, partition_key: str = None):
+        """
+        High-performance bulk load into Actian Vector using vwload.
+
+        Steps:
+        1. Convert Parquet to CSV locally (vwload prefers CSV)
+        2. SCP the CSV to EC2
+        3. Create table schema on Actian
+        4. Run vwload on EC2
+        5. Cleanup remote files
+        """
+        print(f"[Actian] Bulk loading {file_path} into table '{table_name}'...")
+
+        remote_data_dir = self.ssh_params["remote_data_dir"]
+        vwload_path = self.ssh_params["vwload_path"]
+        database = self.connection_params["database"]
+
+        # 1. Read schema and convert to CSV
+        df = pl.read_parquet(file_path)
+
+        with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp:
+            local_csv_path = tmp.name
+            df.write_csv(local_csv_path, include_header=True)
+
+        try:
+            # 2. SCP to EC2
+            remote_csv_path = f"{remote_data_dir}/{table_name}.csv"
+            self._scp_to_remote(local_csv_path, remote_csv_path)
+
+            # 3. Create table schema
+            self._create_table_schema(table_name, df)
+
+            # 4. Run vwload
+            # vwload syntax: vwload [options] dbname tablename filename
+            vwload_cmd = (
+                f"{vwload_path} "
+                f"--fdelim ',' "
+                f"--header "
+                f"--log /tmp/vwload_{table_name}.log "
+                f"{database} {table_name} {remote_csv_path}"
+            )
+
+            print(f"[Actian] Running vwload: {vwload_cmd}")
+            stdout, stderr, exit_code = self._ssh_exec(vwload_cmd, timeout=600)
+
+            if exit_code != 0:
+                # Try to get the log file for more details
+                log_stdout, _, _ = self._ssh_exec(f"cat /tmp/vwload_{table_name}.log 2>/dev/null || echo 'No log'")
+                raise RuntimeError(
+                    f"vwload failed (exit {exit_code}):\n"
+                    f"stderr: {stderr}\n"
+                    f"log: {log_stdout}"
+                )
+
+            print(f"[Actian] Loaded {df.height} rows into {table_name}")
+
+            # 5. Optimize table for analytics
+            self._optimize_table(table_name)
+
+            # 6. Cleanup remote CSV
+            self._ssh_exec(f"rm -f {remote_csv_path}")
+
+        finally:
+            # Cleanup local temp file
+            if os.path.exists(local_csv_path):
+                os.remove(local_csv_path)
+
+    def _create_table_schema(self, table_name: str, df: pl.DataFrame):
+        """Create table schema in Actian based on Polars DataFrame."""
+
+        # Map Polars types to Actian Vector types
+        type_map = {
+            pl.Int8: "TINYINT",
+            pl.Int16: "SMALLINT",
+            pl.Int32: "INTEGER",
+            pl.Int64: "BIGINT",
+            pl.UInt8: "SMALLINT",
+            pl.UInt16: "INTEGER",
+            pl.UInt32: "BIGINT",
+            pl.UInt64: "BIGINT",
+            pl.Float32: "FLOAT4",
+            pl.Float64: "FLOAT8",
+            pl.String: "VARCHAR(1000)",
+            pl.Utf8: "VARCHAR(1000)",
+            pl.Boolean: "BOOLEAN",
+            pl.Date: "DATE",
+            pl.Datetime: "TIMESTAMP",
+        }
+
+        columns = []
+        for col_name, dtype in df.schema.items():
+            # Get base type, handle parameterized types like Datetime
+            actian_type = type_map.get(type(dtype), None)
+            if actian_type is None:
+                actian_type = type_map.get(dtype, "VARCHAR(1000)")
+
+            # Handle datetime subtypes
+            if isinstance(dtype, pl.Datetime):
+                actian_type = "TIMESTAMP"
+
+            columns.append(f'"{col_name}" {actian_type}')
+
+        columns_sql = ", ".join(columns)
+
+        # Drop and recreate for clean state
+        drop_sql = f"DROP TABLE IF EXISTS {table_name}"
+        create_sql = f"CREATE TABLE {table_name} ({columns_sql})"
+
+        self._execute_sql(drop_sql)
+        self._execute_sql(create_sql)
+
+        print(f"[Actian] Created table {table_name} with {len(columns)} columns")
+
+    def _optimize_table(self, table_name: str):
+        """Optimize table for analytical queries."""
+        # MODIFY TO COMBINE reorganizes data for better scan performance
+        try:
+            self._execute_sql(f"MODIFY {table_name} TO COMBINE")
+            print(f"[Actian] Optimized table {table_name}")
+        except Exception as e:
+            print(f"[Actian] Warning: MODIFY TO COMBINE failed (non-fatal): {e}")
+
+    def _execute_sql(self, sql: str):
+        """Execute a SQL statement (DDL/utility) without timing."""
+        database = self.connection_params["database"]
+        sql_path = self.ssh_params["actian_sql_path"]
+
+        escaped_sql = sql.replace("'", "'\\''")
+        command = f"echo '{escaped_sql}' | {sql_path} {database}"
+
+        stdout, stderr, exit_code = self._ssh_exec(command, timeout=120)
+
+        if exit_code != 0:
+            raise RuntimeError(f"SQL execution failed: {stderr}\nSQL: {sql}")
+
+    def close(self):
+        """Close SSH connection."""
+        if self._ssh_client:
+            self._ssh_client.close()
+            self._ssh_client = None

sql_benchmarks/resources/base_engine.py

@@ -0,0 +1,58 @@
+# sql_benchmarks_dagster/resources/base_engine.py
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, Protocol, runtime_checkable
+
+@runtime_checkable
+class IBenchmarkEngine(Protocol):
+    """
+    ABSTRACT BASE CLASS: Defines the mandatory interface for all 
+    pluggable benchmark engines (SQL, Graph, etc.).
+    
+    The benchmark and ingestion factories rely ONLY on this contract.
+    """
+
+    @abstractmethod
+    def get_engine_name(self) -> str:
+        """Returns the canonical name of the engine (e.g., 'duckdb', 'postgres', 'neo4j')."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def clear_cache(self):
+        """
+        Resets the database state to ensure a 'Cold Cache' benchmark.
+        """
+        pass   
+     
+    @abstractmethod
+    def run_query(self,
+                  sql: str,
+                  partition_key: str,
+                  engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """
+        Executes the query for benchmarking and forces result collection.
+
+        Args:
+            sql (str): The SQL or Cypher query to run.
+            partition_key (str): The unique key for the scenario/database instance.
+            engine_params (Dict): THIS engine's own tuning namespace, already
+                                  sliced by the benchmark factory from the
+                                  config's namespaced engine_params block
+                                  (e.g. the postgres engine receives
+                                  {'work_mem': '64MB'}). Each engine owns its
+                                  vocabulary; engines without tunables accept
+                                  the kwarg and ignore it.
+
+        Returns:
+            Optional[float]: The execution duration in seconds, if measured internally.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def bulk_load(self, filepath: str, target_table_name: str, partition_key: str) -> None:
+        """
+        Handles the engine-specific logic for efficiently loading a file (e.g., Parquet).
+        
+        Note: The partition_key is required to map the ingestion to the correct physical 
+              database file (e.g., benchmark_tiny_ssd.duckdb).
+        """
+        raise NotImplementedError

sql_benchmarks/resources/base_schema_client.py

@@ -0,0 +1,104 @@
+"""
+Abstract base for schema-first (non-relational) database clients.
+
+SQL clients (DuckDB, Postgres, Actian) bulk-load via a single file-level
+operation (COPY / read_parquet / vwload).  Schema-first databases require a
+different three-step sequence:
+
+    1. Prepare the data store  (drop / create database, collection, graph …)
+    2. Define the schema       (entity types, attributes, constraints …)
+    3. Insert rows in batches  (no native bulk-file path)
+
+This class owns that orchestration via the Template Method pattern.
+Concrete subclasses implement the three abstract steps for their specific
+database.  A future Neo4j, ArangoDB, or Weaviate client would extend this
+class and implement only those three methods.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+import polars as pl
+
+
+class SchemaFirstClient(ABC):
+    """
+    Template-method base for schema-first database clients.
+
+    Subclass contract
+    -----------------
+    Implement the three abstract methods below.  Do NOT override
+    ``bulk_load`` unless you genuinely need to change the loading sequence.
+    """
+
+    #: Rows per write transaction.  Subclasses may override.
+    BATCH_SIZE: int = 200
+
+    # ------------------------------------------------------------------
+    # Public interface
+    # ------------------------------------------------------------------
+
+    def bulk_load(self, filepath: str, entity_type: str) -> None:
+        """
+        Load a Parquet file into the database.
+
+        Sequence (fixed):
+            1. _prepare_store   — wipe and recreate the data store
+            2. _define_schema   — declare types / structure
+            3. _insert_batch    — write rows in BATCH_SIZE chunks
+        """
+        df = pl.read_parquet(filepath)
+
+        self._prepare_store(entity_type)
+        self._define_schema(entity_type, df)
+
+        rows = df.to_dicts()
+        for i in range(0, len(rows), self.BATCH_SIZE):
+            self._insert_batch(entity_type, rows[i : i + self.BATCH_SIZE])
+
+        print(
+            f"[{type(self).__name__}] Loaded {len(rows)} rows"
+            f" into '{entity_type}'"
+        )
+
+    @abstractmethod
+    def run_query(self, query: str, scenario_params: Dict[str, Any]) -> Optional[float]:
+        """
+        Execute a query and return the wall-clock duration in seconds.
+
+        Implementations must force full result materialisation before
+        stopping the clock (consistent with how DuckDB and Actian are timed).
+        """
+        raise NotImplementedError
+
+    # ------------------------------------------------------------------
+    # Abstract steps (implement in subclass)
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    def _prepare_store(self, entity_type: str) -> None:
+        """
+        Drop any existing data store and create a fresh one.
+
+        Examples:
+            TypeDB  → drop + create database
+            Neo4j   → drop + create named graph
+            MongoDB → drop + create collection
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _define_schema(self, entity_type: str, df: pl.DataFrame) -> None:
+        """
+        Declare the schema / structure for ``entity_type``.
+
+        ``df`` is provided so implementations can derive the schema from
+        the DataFrame's column names and dtypes rather than requiring a
+        separate schema file.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _insert_batch(self, entity_type: str, rows: list) -> None:
+        """Persist one batch of plain-dict rows to the data store."""
+        raise NotImplementedError

sql_benchmarks/resources/duckdb.py

@@ -0,0 +1,58 @@
+import os
+from dagster import ConfigurableResource
+from typing import Dict, Any, Optional
+from .base_engine import IBenchmarkEngine 
+from .duckdb_client import DuckDBClient 
+from pydantic import ConfigDict
+from ..utils.system import thrash_os_cache
+
+# The Engine is the immutable facade that satisfies the contract.
+class DuckDBEngine(ConfigurableResource): 
+    # --- CONFIGURATION (Immutable) ---
+    data_folder: str # This is the only configuration it holds
+    model_config = ConfigDict(extra='forbid')
+    # --- FACTORY METHOD ---
+    # This is the Factory Method: It creates the worker/client instance.
+    def _get_client(self) -> 'DuckDBClient':
+        """Instantiates the stateful client for execution."""
+        return DuckDBClient(data_folder=self.data_folder)
+    
+    # --- IBenchmarkEngine Implementation (Delegation) ---
+
+    def get_engine_name(self) -> str:
+        return "duckdb"
+
+    def bulk_load(self, filepath: str, table_name: str, partition_key: str, table_def: dict = None) -> None:
+        """Delegates the bulk loading operation to the client."""
+        client = self._get_client()
+        client.bulk_load(filepath, table_name, partition_key, table_def)
+
+             
+    def run_query(self,
+                  sql: str,
+                  partition_key: str,
+                  engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """Delegates the core benchmarking query execution to the client."""
+        self.clear_cache()
+
+        client = self._get_client()
+        # Delegate the call using the exact method signature
+        return client.run_query(sql=sql,
+                                partition_key=partition_key,
+                                engine_params=engine_params
+                                )
+
+    # --- Utility Methods ---
+    
+    def _get_db_path(self, partition_key: str):
+        """Calculates the database file path using the partition_key (still required for path lookups)."""
+        # NOTE: This method can safely remain here as it only involves string manipulation, not I/O.
+        if partition_key is None:
+            return os.path.join(self.data_folder, "benchmark.duckdb") 
+        
+        db_filename = f"benchmark_{partition_key}.duckdb" 
+        return os.path.join(self.data_folder, db_filename)
+
+    def clear_cache(self):
+        
+        thrash_os_cache()

sql_benchmarks/resources/duckdb_client.py

@@ -0,0 +1,157 @@
+import duckdb
+import os
+import re
+import time
+from contextlib import contextmanager
+from typing import Dict, Any, Optional
+from ..utils.system import thrash_os_cache
+from .base_engine import IBenchmarkEngine
+
+_SAFE_IDENTIFIER = re.compile(r"^\w+$")  # letters, digits, underscores only
+
+
+def _assert_safe_identifier(value: str, label: str) -> None:
+    if not _SAFE_IDENTIFIER.match(value):
+        raise ValueError(f"Unsafe {label}: '{value}'. Only word characters allowed.")
+
+
+# The 'duckdb' engine_params namespace vocabulary: session settings applied
+# via SET before the measured query. Strict allowlist — an unknown key is a
+# config error and must fail loudly, never be silently ignored.
+DUCKDB_SETTING_KEYS = frozenset({
+    "threads",
+    "memory_limit",
+})
+
+# SET cannot use parameter binding; values are interpolated as quoted
+# literals, so restrict the alphabet (covers ints and sizes like '1GB').
+_SAFE_SETTING_VALUE = re.compile(r"^[A-Za-z0-9._]+$")
+
+
+def _apply_engine_params(con, engine_params: Dict[str, Any]) -> None:
+    for key, value in (engine_params or {}).items():
+        if key not in DUCKDB_SETTING_KEYS:
+            raise ValueError(
+                f"Unknown duckdb engine_params key '{key}'. "
+                f"Allowed: {sorted(DUCKDB_SETTING_KEYS)}"
+            )
+        if not _SAFE_SETTING_VALUE.match(str(value)):
+            raise ValueError(f"Unsafe value for duckdb setting '{key}': {value!r}")
+        con.execute(f"SET {key} = '{value}'")
+
+# Note: DuckDBClient does NOT inherit ConfigurableResource
+
+# ---------------------------------------------------------------------------
+# Module-level db-file tracker (mirrors the TypeDBEngine pattern).
+# The first bulk_load call for a given db file deletes any stale copy
+# (clean slate).  Subsequent tables are added to the same file via
+# CREATE OR REPLACE TABLE, so multi-table experiments (e.g. hypergraph
+# supply chain) work correctly.  Keyed by db_path — not partition_key —
+# because multiple engines (duckdb, quack) load the same partitions into
+# different data folders.  Cleared naturally between runs because each
+# execute_run.py invocation starts a fresh Python process.
+# ---------------------------------------------------------------------------
+_INITIALIZED_DB_FILES: set = set()
+
+class DuckDBClient: 
+    """
+    STATEFUL CLIENT: Holds the execution logic and performs all I/O operations,
+    ensuring the Engine resource remains immutable.
+    """
+    
+    def __init__(self, data_folder: str):
+        # Configuration required for execution, held on the stateful object
+        self.data_folder = data_folder
+    
+    # --- IBenchmarkEngine Implementation (Execution Logic) ---
+    
+    def bulk_load(self, filepath: str, target_table_name: str, partition_key: str, table_def: dict = None) -> None:
+        # table_def is accepted for interface symmetry; DuckDB is columnar and
+        # relies on automatic min-max zonemaps, so declared B-tree indexes/PKs
+        # are not applied here.
+        db_path = self._get_db_path(partition_key)
+
+        # Validate inputs before constructing SQL.
+        # Table name: DuckDB doesn't support parameterized identifiers, so we validate
+        # with a strict allowlist regex (word chars only) before interpolation.
+        # Filepath: use read_parquet(?) parameterized binding for the value.
+        _assert_safe_identifier(target_table_name, "table name")
+        filepath = os.path.realpath(filepath)  # resolve symlinks / traversal
+
+        # First table for this db file: wipe any stale copy for a clean
+        # slate.  Subsequent tables are added to the same file via
+        # CREATE OR REPLACE TABLE so all tables coexist.
+        if db_path not in _INITIALIZED_DB_FILES:
+            if os.path.exists(db_path):
+                try:
+                    os.remove(db_path)
+                except OSError as e:
+                    print(f"[WARN] Could not remove existing DB file {db_path}: {e}")
+            _INITIALIZED_DB_FILES.add(db_path)
+
+        os.makedirs(os.path.dirname(db_path), exist_ok=True)
+        with duckdb.connect(db_path, read_only=False) as con:
+            con.execute(
+                f"CREATE OR REPLACE TABLE {target_table_name} AS SELECT * FROM read_parquet(?)",
+                [filepath],
+            )
+             
+    def run_query(self,
+                  sql: str,
+                  partition_key: str,
+                  engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """Execute a benchmark query against the partition's DuckDB file.
+
+        engine_params is the 'duckdb' namespace (threads, memory_limit),
+        applied via SET before timing starts so the measurement reflects the
+        configured execution, not the configuration itself.
+        """
+        db_path = self._get_db_path(partition_key)
+
+        # Enable RW for Sentinel experiments (CREATE TABLE)
+        with duckdb.connect(db_path, read_only=False) as con:
+            _apply_engine_params(con, engine_params)
+            start = time.time()
+            # In DuckDB, multiple statements in one string are executed if separated by semicolon
+            # .sql() executes them. 
+            con.sql(sql).fetchall()
+            end = time.time()
+            
+            return end - start
+            
+    # --- Internal/Utility Methods (Moved here) ---
+    
+    def _get_db_path(self, partition_key: str):
+        """Calculates the database file path using the partition_key."""
+        if partition_key is None:
+            return os.path.join(self.data_folder, "benchmark.duckdb") 
+        
+        db_filename = f"benchmark_{partition_key}.duckdb" 
+        return os.path.join(self.data_folder, db_filename)
+
+    # --- Utility Methods ---
+    # execute_query, get_connection, and execute_on_file are moved here
+    # as they all involve I/O and state.
+
+    def execute_query(self, sql: str, partition_key: str = None, read_only: bool = False, is_benchmark: bool = False):
+        db_path = self._get_db_path(partition_key)
+        
+        if not read_only:
+            os.makedirs(os.path.dirname(db_path), exist_ok=True)
+        
+        with duckdb.connect(db_path, read_only=read_only) as con:
+            result = con.execute(sql)
+            if is_benchmark:
+                result.fetchall()
+    
+    @contextmanager
+    def get_connection(self, db_path: str, read_only: bool = False):
+        con = duckdb.connect(db_path, read_only=read_only)
+        try:
+            yield con
+        finally:
+            con.close()
+
+    def execute_on_file(self, sql: str, target_path: str):
+        with self.get_connection(target_path) as con:
+            con.execute(sql)