diff --git a/CHANGELOG.md b/CHANGELOG.md index a559c418a0..6956b1961b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -310,6 +310,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 there. Receipt entries for constraints carry `kind: "constraint"` and a `constraint_type`. Quoted non-ASCII table names, such as "É", are now found consistently with DuckDB's `lower(table_name)` catalog filters. +- **Tuning statements issued through `.sql()` and `.query()` are recorded.** + Spark, PySpark, LakeSail, Velox, DataFusion and BigQuery tuned schema DDL now + reaches the applied-tuning ledger, so a tuned Spark run no longer reports + `noop`. BigQuery statements are recorded with the final state of the job, not + at submission. Multi-statement strings are split before recording, so a + leading `SET` can no longer hide a tuning statement. Schema capture also + recognises `PARTITIONED BY`, `CLUSTERED BY`, `DISTKEY`, `DISTSTYLE`, + `DUPLICATE KEY`, `PRIMARY INDEX`, `distribution =`, `partitioning =`, + `sorted_by` and `bucketed_by`. + Spark table definitions now put `USING ` before a trailing + `PARTITIONED BY` or `CLUSTER BY`, which Spark requires. ### Changed diff --git a/benchbox/core/tuning/applied_ledger.py b/benchbox/core/tuning/applied_ledger.py index 5725242144..225d08a4ad 100644 --- a/benchbox/core/tuning/applied_ledger.py +++ b/benchbox/core/tuning/applied_ledger.py @@ -66,11 +66,26 @@ r"^\s*(?:create\s+index|create\s+(?:or\s+replace\s+)?table|alter\s+table)\b", re.IGNORECASE, ) -_SCHEMA_TUNING_FOOTPRINT_RE = re.compile( - r"\b(?:cluster\s+by|distributed\s+by|foreign\s+key|order\s+by|partition\s+by|primary\s+key|" - r"sortkey|unique|check)\b", +_SCHEMA_TUNING_PHRASE_RE = re.compile( + r"\b(?:cluster\s+by|clustered\s+by|distributed\s+by|foreign\s+key|order\s+by|partition\s+by|" + r"partitioned\s+by|primary\s+key|primary\s+index|duplicate\s+key)\b" + r"|\b(?:distribution|partitioning)\s*=", re.IGNORECASE, ) +_SCHEMA_TUNING_KEYWORD_RE = re.compile( + r"\b(?:sortkey|distkey|diststyle|sorted_by|bucketed_by|unique|check)\b", + re.IGNORECASE, +) +_UNIQUE_NULLS_DISTINCT_RE = re.compile(r"\s*nulls\s+(?:not\s+distinct|distinct)\s*\(", re.IGNORECASE) + + +def _bare_keyword_is_clause(shape: str, match: re.Match[str]) -> bool: + rest = shape[match.end() :] + if rest.lstrip()[:1] in ("(", "="): + return True + if match.group(0).lower() == "unique" and _UNIQUE_NULLS_DISTINCT_RE.match(rest) is not None: + return True + return not shape[: match.start()].rstrip().endswith(("(", ",")) def _sql_shape(statement: Any) -> str: @@ -85,7 +100,9 @@ def is_schema_tuning_statement(statement: Any) -> bool: return False if re.match(r"^\s*create\s+index\b", shape, re.IGNORECASE): return True - return _SCHEMA_TUNING_FOOTPRINT_RE.search(shape) is not None + if _SCHEMA_TUNING_PHRASE_RE.search(shape) is not None: + return True + return any(_bare_keyword_is_clause(shape, match) for match in _SCHEMA_TUNING_KEYWORD_RE.finditer(shape)) def _split_sql_script(script: Any) -> list[str]: @@ -97,16 +114,18 @@ def _split_sql_script(script: Any) -> list[str]: import sqlite3 statements: list[str] = [] - pending: list[str] = [] - for line in text.splitlines(keepends=True): - pending.append(line) - candidate = "".join(pending) + start = 0 + for index, char in enumerate(text): + if char != ";": + continue + candidate = text[start : index + 1] if sqlite3.complete_statement(candidate): if candidate.strip(): statements.append(candidate.strip()) - pending.clear() - if pending and "".join(pending).strip(): - statements.append("".join(pending).strip()) + start = index + 1 + tail = text[start:].strip() + if tail: + statements.append(tail) return statements except Exception: # pragma: no cover return [text] @@ -127,10 +146,13 @@ def normalize_ledger_phase(phase: Any) -> str: def _is_recordable_statement(statement: Any) -> bool: try: - text = str(statement).lstrip().lstrip("(").lstrip().lower() + text = _SQL_COMMENT_RE.sub(" ", str(statement)) except Exception: # pragma: no cover return True - return not text.startswith(_READBACK_PREFIXES) + stripped = text.strip().strip(";").strip() + if not stripped: + return False + return not stripped.lstrip("(").lstrip().lower().startswith(_READBACK_PREFIXES) @dataclass @@ -322,9 +344,178 @@ def _should_record(self, statement: Any) -> bool: logger.debug("applied-ledger statement filter degraded: %s", exc) return True + def _recordable_statements(self, statement: Any) -> list[str]: + try: + pieces = _split_sql_script(statement) + if len(pieces) <= 1: + pieces = [statement] + return [str(piece) for piece in pieces if self._should_record(piece)] + except Exception as exc: + logger.debug("applied-ledger statement capture degraded: %s", exc) + return [] + + def _record_all(self, statements: list[str], status: str, error: Any | None = None) -> None: + for statement in statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def _run(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: + statements = self._recordable_statements(statement) + try: + result = fn(statement, *args, **kwargs) + except Exception as exc: + self._record_all(statements, STATEMENT_FAILED, exc) + raise + self._record_all(statements, EXECUTED) + return result + + def _run_verb(self, fn: Any, args: tuple, kwargs: dict) -> Any: + statements: list[str] = [] + try: + name = _first_parameter_name(fn) + if args: + statement, rest, rest_kwargs = args[0], args[1:], kwargs + elif name in kwargs: + statement, rest, rest_kwargs = kwargs[name], args, {k: v for k, v in kwargs.items() if k != name} + else: + raise LookupError(name) + if not _is_dry_run(rest, rest_kwargs): + statements = self._recordable_statements(statement) + except Exception as exc: + logger.debug("applied-ledger statement capture degraded: %s", exc) + try: + result = fn(*args, **kwargs) + except Exception as exc: + self._record_all(statements, STATEMENT_FAILED, exc) + raise + if not statements: + return result + if callable(getattr(result, "result", None)) and callable(getattr(result, "done", None)): + return _RecordingJob(result, self._ledger, self._phase, statements) + if _is_lazy_frame(result): + return _RecordingFrame(result, self._ledger, self._phase, statements) + self._record_all(statements, EXECUTED) + return result + + +def _first_parameter_name(fn: Any) -> str | None: + import inspect + + for parameter in inspect.signature(fn).parameters.values(): + if parameter.kind in (parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY): + return parameter.name + break + return None + + +def _is_dry_run(args: tuple, kwargs: dict) -> bool: + config = kwargs.get("job_config") or (args[0] if args else None) + return getattr(config, "dry_run", False) is True + + +class _RecordingJob: + __slots__ = ("_job", "_ledger", "_phase", "_statements", "_settled") + + def __init__(self, job: Any, ledger: AppliedTuningLedger, phase: str, statements: list[str]) -> None: + object.__setattr__(self, "_job", job) + object.__setattr__(self, "_ledger", ledger) + object.__setattr__(self, "_phase", phase) + object.__setattr__(self, "_statements", statements) + object.__setattr__(self, "_settled", False) + + def _settle(self, status: str, error: Any | None = None) -> None: + if self._settled: + return + object.__setattr__(self, "_settled", True) + for statement in self._statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def result(self, *args: Any, **kwargs: Any) -> Any: + from concurrent.futures import TimeoutError as FutureTimeoutError + + try: + outcome = self._job.result(*args, **kwargs) + except FutureTimeoutError: + raise + except Exception as exc: + self._settle(STATEMENT_FAILED, exc) + raise + self._settle(EXECUTED) + return outcome + + def exception(self, *args: Any, **kwargs: Any) -> Any: + error = self._job.exception(*args, **kwargs) + if error is not None: + self._settle(STATEMENT_FAILED, error) + elif self._job.done(): + self._settle(EXECUTED) + return error + + def __getattr__(self, name: str) -> Any: + return getattr(self._job, name) + + +_LAZY_FRAME_SUBMISSION_METHODS = frozenset({"collect", "show", "fetchall", "fetch_all"}) +_LAZY_FRAME_ACTION_METHODS = frozenset( + { + "collect", + "show", + "fetchall", + "fetch_all", + "count", + "to_pandas", + "toPandas", + "to_polars", + "to_arrow", + } +) + + +def _is_lazy_frame(result: Any) -> bool: + try: + for name in _LAZY_FRAME_SUBMISSION_METHODS: + if callable(getattr(result, name, None)): + return True + except Exception: + return False + return False + + +class _RecordingFrame: + __slots__ = ("_frame", "_ledger", "_phase", "_statements", "_settled") + + def __init__(self, frame: Any, ledger: AppliedTuningLedger, phase: str, statements: list[str]) -> None: + object.__setattr__(self, "_frame", frame) + object.__setattr__(self, "_ledger", ledger) + object.__setattr__(self, "_phase", phase) + object.__setattr__(self, "_statements", statements) + object.__setattr__(self, "_settled", False) + + def _settle(self, status: str, error: Any | None = None) -> None: + if self._settled: + return + object.__setattr__(self, "_settled", True) + for statement in self._statements: + self._ledger.record(statement, self._phase, status=status, error=error) + + def __getattr__(self, name: str) -> Any: + attribute = getattr(self._frame, name) + if name in _LAZY_FRAME_ACTION_METHODS and callable(attribute): + + def call(*args: Any, **kwargs: Any) -> Any: + try: + outcome = attribute(*args, **kwargs) + except Exception as exc: + self._settle(STATEMENT_FAILED, exc) + raise + self._settle(EXECUTED) + return outcome + + return call + return attribute + class RecordingConnection(_RecordingProxy): - __slots__ = ("_conn",) + __slots__ = ("_conn", "_verbs") def __init__( self, @@ -332,9 +523,11 @@ def __init__( ledger: AppliedTuningLedger, phase: str, statement_filter: Callable[[Any], bool] | None = None, + execute_verbs: tuple[str, ...] = (), ) -> None: super().__init__(ledger, phase, statement_filter) object.__setattr__(self, "_conn", connection) + object.__setattr__(self, "_verbs", frozenset(execute_verbs)) def execute(self, statement: Any, *args: Any, **kwargs: Any) -> Any: return self._run(self._conn.execute, statement, args, kwargs) @@ -343,26 +536,27 @@ def cursor(self, *args: Any, **kwargs: Any) -> Any: return _RecordingCursor(self._conn.cursor(*args, **kwargs), self._ledger, self._phase, self._statement_filter) def executescript(self, script: Any, *args: Any, **kwargs: Any) -> Any: - result = self._conn.executescript(script, *args, **kwargs) - for statement in _split_sql_script(script): - if self._should_record(statement): - self._ledger.record(statement, self._phase, status=EXECUTED) - return result - - def _run(self, fn: Any, statement: Any, args: tuple, kwargs: dict) -> Any: - record = self._should_record(statement) try: - result = fn(statement, *args, **kwargs) + result = self._conn.executescript(script, *args, **kwargs) except Exception as exc: - if record: - self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) + for statement in _split_sql_script(script): + if self._should_record(statement): + self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) raise - if record: - self._ledger.record(statement, self._phase, status=EXECUTED) + for statement in _split_sql_script(script): + if self._should_record(statement): + self._ledger.record(statement, self._phase, status=EXECUTED) return result def __getattr__(self, name: str) -> Any: - return getattr(self._conn, name) + attribute = getattr(self._conn, name) + if name in self._verbs and callable(attribute): + + def call(*args: Any, **kwargs: Any) -> Any: + return self._run_verb(attribute, args, kwargs) + + return call + return attribute def __setattr__(self, name: str, value: Any) -> None: setattr(self._conn, name, value) @@ -393,16 +587,7 @@ def __init__( object.__setattr__(self, "_cur", cursor) def execute(self, statement: Any, *args: Any, **kwargs: Any) -> Any: - record = self._should_record(statement) - try: - result = self._cur.execute(statement, *args, **kwargs) - except Exception as exc: - if record: - self._ledger.record(statement, self._phase, status=STATEMENT_FAILED, error=exc) - raise - if record: - self._ledger.record(statement, self._phase, status=EXECUTED) - return result + return self._run(self._cur.execute, statement, args, kwargs) def __getattr__(self, name: str) -> Any: return getattr(self._cur, name) @@ -426,11 +611,12 @@ def recording_connection( ledger: AppliedTuningLedger | None, phase: str, statement_filter: Callable[[Any], bool] | None = None, + execute_verbs: tuple[str, ...] = (), ) -> Any: if ledger is None: return connection try: - return RecordingConnection(connection, ledger, phase, statement_filter) + return RecordingConnection(connection, ledger, phase, statement_filter, execute_verbs) except Exception as exc: # pragma: no cover logger.debug("applied-ledger connection wrap degraded: %s", exc) return connection diff --git a/benchbox/platforms/_spark_helpers.py b/benchbox/platforms/_spark_helpers.py index 7868cd0df9..dfd15de49a 100644 --- a/benchbox/platforms/_spark_helpers.py +++ b/benchbox/platforms/_spark_helpers.py @@ -157,6 +157,42 @@ def get_spark_query_plan( _LEADING_COMMENTS_RE = re.compile(r"\A(?:\s*(?:--[^\n]*(?:\n|$)|/\*.*?\*/))+\s*", re.DOTALL) _FIXED_SIZE_ARRAY_TYPE_RE = re.compile(r"\b([A-Za-z]\w*)\s*\[\s*\d+\s*\]") _ARRAY_FIXED_SIZE_SUFFIX_RE = re.compile(r"(>)\s*\[\s*\d+\s*\]") +_CREATE_TABLE_COLUMN_LIST_RE = re.compile( + r"\A\s*CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(?:\"[^\"]*\"|`[^`]*`|[^\s(\"`])+\s*\(", + re.IGNORECASE, +) + + +def _column_list_end(statement: str) -> int: + match = _CREATE_TABLE_COLUMN_LIST_RE.match(statement) + if not match: + return len(statement) + depth = 1 + quote = "" + cursor = match.end() + while cursor < len(statement): + ch = statement[cursor] + cursor += 1 + if quote: + if ch == "\\" and quote == "'": + cursor += 1 + elif ch == quote: + quote = "" + elif ch in "'\"`": + quote = ch + elif ch == "-" and statement[cursor : cursor + 1] == "-": + newline = statement.find("\n", cursor + 1) + cursor = len(statement) if newline == -1 else newline + 1 + elif ch == "/" and statement[cursor : cursor + 1] == "*": + close = statement.find("*/", cursor + 1) + cursor = len(statement) if close == -1 else close + 2 + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + return cursor + return len(statement) def _strip_balanced_paren_constraints(statement: str) -> str: @@ -227,7 +263,9 @@ def optimize_spark_table_definition( if ")" in statement: fmt = (table_format or "parquet").upper() - statement = statement.rstrip(";").rstrip() + f" USING {fmt}" + statement = statement.rstrip(";").rstrip() + end = _column_list_end(statement) + statement = f"{statement[:end]} USING {fmt}{statement[end:]}" return statement @@ -255,6 +293,8 @@ def purge_orphaned_warehouse_directory(spark: Any, *, logger: logging.Logger) -> class SparkLikeAdapterMixin: + ledger_execute_verbs = ("sql",) + def apply_constraint_configuration( self, primary_key_config: Any, @@ -320,6 +360,13 @@ def _record_spark_conf_ledger(ledger: Any, key: str, value: Any, *, applied: boo pass +def _run_schema_statement(spark: Any, statement: str) -> None: + result = spark.sql(statement) + collect = getattr(result, "collect", None) + if callable(collect): + collect() + + def run_spark_schema_creation_loop( spark: Any, statements: list[str], @@ -341,7 +388,7 @@ def run_spark_schema_creation_loop( if not statement.strip(): continue try: - spark.sql(statement) + _run_schema_statement(spark, statement) logger.debug(f"Executed schema statement: {statement[:100]}...") except Exception as exc: error_lower = str(exc).lower() @@ -359,7 +406,7 @@ def run_spark_schema_creation_loop( "drop the conflicting object manually." ) from exc - spark.sql(f"DROP TABLE IF EXISTS {table_name}") + _run_schema_statement(spark, f"DROP TABLE IF EXISTS {table_name}") if on_location_collision is not None and "location_already_exists" in error_lower: on_location_collision(spark, table_name) - spark.sql(statement) + _run_schema_statement(spark, statement) diff --git a/benchbox/platforms/base/adapter.py b/benchbox/platforms/base/adapter.py index 6b198f504f..9465461681 100644 --- a/benchbox/platforms/base/adapter.py +++ b/benchbox/platforms/base/adapter.py @@ -594,7 +594,12 @@ def run_enhanced_benchmark(self, benchmark, **run_config) -> EnhancedBenchmarkRe benchmark_type = run_config.get("benchmark_type", "olap") self.configure_for_benchmark( - recording_connection(connection, self._applied_tuning_ledger, PHASE_SESSION), + recording_connection( + connection, + self._applied_tuning_ledger, + PHASE_SESSION, + execute_verbs=self.ledger_execute_verbs, + ), benchmark_type, ) @@ -921,6 +926,7 @@ def _setup_fresh_database_phases(self, benchmark, connection: Any, effective_tun getattr(self, "_applied_tuning_ledger", None), PHASE_DDL, statement_filter=is_schema_tuning_statement, + execute_verbs=self.ledger_execute_verbs, ) if schema_records_ddl else connection diff --git a/benchbox/platforms/base/sorted_ingestion.py b/benchbox/platforms/base/sorted_ingestion.py index 606768c8cf..d2eaaf9d56 100644 --- a/benchbox/platforms/base/sorted_ingestion.py +++ b/benchbox/platforms/base/sorted_ingestion.py @@ -184,6 +184,28 @@ def _resolve_ctas_sort_columns(self, table_name: str, tuning_config: Any) -> lis return sort_columns + def _recording_ctas_connection(self, connection: Any) -> Any: + from benchbox.core.tuning.applied_ledger import ( + PHASE_POST_LOAD, + RecordingConnection, + is_schema_tuning_statement, + recording_connection, + ) + + if getattr(self, "_applied_tuning_ledger", None) is None: + return connection + if isinstance(connection, RecordingConnection): + return connection + if not callable(getattr(connection, "execute", None)) and not callable(getattr(connection, "sql", None)): + return connection + return recording_connection( + connection, + self._applied_tuning_ledger, + PHASE_POST_LOAD, + statement_filter=is_schema_tuning_statement, + execute_verbs=("sql",), + ) + def _execute_ctas_sort( self, table_name: str, @@ -206,6 +228,7 @@ def _execute_ctas_sort( return True ledger_statement = "\n".join(statements) + target = self._recording_ctas_connection(connection) transaction_started = False try: if atomic: @@ -214,7 +237,7 @@ def _execute_ctas_sort( self._record_sorted_ingestion_skip(validated_table, sorted_columns, reason) self.logger.info("Skipping constraint-preserving CTAS sort for %s: %s", table_name, reason) return False - self._execute_sql_statement(connection, "BEGIN TRANSACTION;") + self._execute_sql_statement(target, "BEGIN TRANSACTION;") transaction_started = True if atomic and self._has_populated_fk_dependents(connection, validated_table): self._rollback_transaction(connection) @@ -228,13 +251,13 @@ def _execute_ctas_sort( ) return False for statement in statements: - self._execute_sql_statement(connection, statement) + self._execute_sql_statement(target, statement) if atomic: commit_method = getattr(connection, "commit", None) if callable(commit_method): commit_method() else: - self._execute_sql_statement(connection, "COMMIT;") + self._execute_sql_statement(target, "COMMIT;") transaction_started = False except Exception as exc: rollback_error: Exception | None = None diff --git a/benchbox/platforms/base/tuning.py b/benchbox/platforms/base/tuning.py index b116d248dd..23b616220d 100644 --- a/benchbox/platforms/base/tuning.py +++ b/benchbox/platforms/base/tuning.py @@ -62,6 +62,7 @@ class TuningHooksMixin: _supported_tuning_type_names: Iterable[str] | None = None physical_identifier_case: str = "preserve" post_load_connection_recording: bool = True + ledger_execute_verbs: tuple[str, ...] = () @staticmethod def table_tuning_for(effective_config: Any, table_name: str) -> Any: @@ -85,7 +86,11 @@ def run_post_load_tunings(self, table_name: str, effective_config: Any, connecti ledger = getattr(self, "_applied_tuning_ledger", None) record = ledger is not None and self.post_load_connection_recording - target = recording_connection(connection, ledger, PHASE_POST_LOAD) if record else connection + target = ( + recording_connection(connection, ledger, PHASE_POST_LOAD, execute_verbs=self.ledger_execute_verbs) + if record + else connection + ) performed = False start = mono_time() try: diff --git a/benchbox/platforms/base/tuning_config.py b/benchbox/platforms/base/tuning_config.py index 3b84d27265..84dce35299 100644 --- a/benchbox/platforms/base/tuning_config.py +++ b/benchbox/platforms/base/tuning_config.py @@ -13,7 +13,12 @@ def apply_standard_unified_tuning(adapter: Any, unified_config: UnifiedTuningCon from benchbox.core.tuning.applied_ledger import PHASE_DDL, recording_connection - recording = recording_connection(connection, getattr(adapter, "_applied_tuning_ledger", None), PHASE_DDL) + recording = recording_connection( + connection, + getattr(adapter, "_applied_tuning_ledger", None), + PHASE_DDL, + execute_verbs=getattr(adapter, "ledger_execute_verbs", ()), + ) adapter.apply_constraint_configuration(unified_config.primary_keys, unified_config.foreign_keys, recording) if unified_config.platform_optimizations: diff --git a/benchbox/platforms/bigquery.py b/benchbox/platforms/bigquery.py index 6daa2d9806..65ae7b4d94 100644 --- a/benchbox/platforms/bigquery.py +++ b/benchbox/platforms/bigquery.py @@ -86,6 +86,7 @@ class BigQueryAdapter(PlatformAdapter): physical_identifier_case = "upper" _DELIMITED_FORMATS = frozenset({"tbl", "csv"}) plan_capture_phase_eligible = False + ledger_execute_verbs = ("query",) def __init__(self, **config): super().__init__(**config) diff --git a/benchbox/platforms/datafusion.py b/benchbox/platforms/datafusion.py index 3c26edd651..cb5d8f1703 100644 --- a/benchbox/platforms/datafusion.py +++ b/benchbox/platforms/datafusion.py @@ -136,6 +136,7 @@ class DataFusionAdapter(NoConstraintEnforcementMixin, PlatformAdapter): driver_isolation_capability = DriverIsolationCapability.SUPPORTED supports_external_tables = True plan_capture_phase_eligible = True + ledger_execute_verbs = ("sql",) _process_working_dir_lock_depth: dict[str, int] = {} _process_working_dir_lock_guard = threading.Lock() diff --git a/tests/integration/platforms/test_spark_ledger_capture_live.py b/tests/integration/platforms/test_spark_ledger_capture_live.py new file mode 100644 index 0000000000..45abd90f23 --- /dev/null +++ b/tests/integration/platforms/test_spark_ledger_capture_live.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import logging +from types import SimpleNamespace + +import pytest + +from benchbox.core.tuning.applied_ledger import ( + EXECUTED, + NOOP, + PHASE_DDL, + AppliedTuningLedger, + is_schema_tuning_statement, + recording_connection, +) + +from .delta_live_helpers import delta_live_skip_reason, make_delta_spark_session + +pytestmark = [ + pytest.mark.integration, + pytest.mark.live_integration, + pytest.mark.live_delta, + pytest.mark.skipif( + delta_live_skip_reason() is not None, + reason=delta_live_skip_reason() or "PySpark + Delta Lake runtime unavailable", + ), +] + + +@pytest.fixture(scope="module") +def spark(tmp_path_factory, pyspark_test_environment): + session = make_delta_spark_session(tmp_path_factory.mktemp("ledger_capture_warehouse"), app_name="ledger-capture") + session.sql("CREATE DATABASE IF NOT EXISTS ledger_capture") + session.sql("USE ledger_capture") + yield session + session.stop() + + +def test_partitioned_and_clustered_ddl_is_recorded_and_status_is_not_noop(spark) -> None: + from benchbox.platforms.spark import SparkAdapter + + ledger = AppliedTuningLedger() + assert ledger.overall_status(tuning_enabled=True, has_config=True) == NOOP + connection = recording_connection( + spark, + ledger, + PHASE_DDL, + statement_filter=is_schema_tuning_statement, + execute_verbs=SparkAdapter.ledger_execute_verbs, + ) + adapter = SimpleNamespace( + _create_schema_with_tuning=lambda *args, **kwargs: ( + "CREATE TABLE orders (o_id INT, o_date DATE)\nPARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT)\nCLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="delta", + logger=logging.getLogger(__name__), + _remove_orphaned_table_location=lambda *args, **kwargs: None, + ) + + SparkAdapter.create_schema(adapter, SimpleNamespace(), connection) + + recorded = [(s.phase, s.status, s.statement) for s in ledger.statements] + assert [(phase, status) for phase, status, _ in recorded] == [(PHASE_DDL, EXECUTED)] * 2 + assert "PARTITIONED BY (o_date)" in recorded[0][2] + assert "CLUSTER BY (l_key)" in recorded[1][2] + assert ledger.overall_status(tuning_enabled=True, has_config=True) != NOOP + assert spark.sql("SELECT COUNT(*) FROM plain").collect()[0][0] == 0 + assert {table.name for table in spark.catalog.listTables()} >= {"orders", "lineitem", "plain"} diff --git a/tests/unit/core/tuning/test_ledger_execute_verbs.py b/tests/unit/core/tuning/test_ledger_execute_verbs.py new file mode 100644 index 0000000000..7e0b046c5a --- /dev/null +++ b/tests/unit/core/tuning/test_ledger_execute_verbs.py @@ -0,0 +1,632 @@ +from __future__ import annotations + +import logging +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from benchbox.core.tuning.applied_ledger import ( + EXECUTED, + PHASE_DDL, + STATEMENT_FAILED, + AppliedTuningLedger, + is_schema_tuning_statement, + recording_connection, +) +from benchbox.core.tuning.introspection import corroborate + +pytestmark = [ + pytest.mark.unit, + pytest.mark.fast, +] + + +class _SqlDriver: + def __init__(self, fail_on: str | None = None) -> None: + self.calls: list[str] = [] + self._fail_on = fail_on + + def sql(self, statement, *args, **kwargs): + if self._fail_on is not None and self._fail_on in str(statement): + raise RuntimeError("driver rejected statement") + self.calls.append(str(statement)) + return None + + def execute(self, statement, *args, **kwargs): + self.calls.append(str(statement)) + return "executed" + + +class _LazyFrame: + def __init__(self, error: Exception | None = None) -> None: + self._error = error + self.actions: list[str] = [] + + def collect(self, *args, **kwargs): + self.actions.append("collect") + if self._error is not None: + raise self._error + return [("row",)] + + def show(self, *args, **kwargs): + self.actions.append("show") + if self._error is not None: + raise self._error + return None + + def count(self, *args, **kwargs): + self.actions.append("count") + if self._error is not None: + raise self._error + return 1 + + +class _LazySqlDriver: + def __init__(self, error: Exception | None = None) -> None: + self.calls: list[str] = [] + self.frames: list[_LazyFrame] = [] + self._error = error + + def sql(self, statement, *args, **kwargs): + self.calls.append(str(statement)) + frame = _LazyFrame(self._error) + self.frames.append(frame) + return frame + + +class _FakeJob: + def __init__(self, error: Exception | None = None) -> None: + self._error = error + self.state = "RUNNING" + self.job_id = "job-1" + + def result(self, *args, **kwargs): + if self._error is not None: + raise self._error + self.state = "DONE" + return "rows" + + def done(self) -> bool: + return self.state == "DONE" + + def exception(self, *args, **kwargs): + return self._error + + +class _QueryClient: + def __init__(self, error: Exception | None = None) -> None: + self.submitted: list[str] = [] + self.jobs: list[_FakeJob] = [] + self._error = error + + def query(self, statement, job_config=None, **kwargs): + self.submitted.append(str(statement)) + job = _FakeJob(self._error) + self.jobs.append(job) + return job + + def dataset(self, dataset_id): + return f"ref::{dataset_id}" + + def get_dataset(self, ref): + return ref + + +def _ddl_proxy(driver, verbs, ledger=None): + ledger = ledger if ledger is not None else AppliedTuningLedger() + return ( + recording_connection( + driver, + ledger, + PHASE_DDL, + statement_filter=is_schema_tuning_statement, + execute_verbs=verbs, + ), + ledger, + ) + + +class TestSqlVerb: + def test_partitioned_and_clustered_ddl_is_recorded_as_ddl(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql("CREATE TABLE t (a INT, b INT) USING delta PARTITIONED BY (a) CLUSTER BY (b)") + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + assert "PARTITIONED BY (a)" in ledger.statements[0].statement + + def test_readback_through_sql_executes_but_is_not_recorded(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + for readback in ("SELECT 1", "SHOW TABLES", "DESCRIBE t", "EXPLAIN SELECT 1"): + proxy.sql(readback) + + assert len(driver.calls) == 4 + assert ledger.is_empty() + + def test_verb_is_not_wrapped_unless_declared(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ()) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert driver.calls == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + assert ledger.is_empty() + + def test_failed_statement_is_recorded_then_reraised(self) -> None: + driver = _SqlDriver(fail_on="PARTITIONED") + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + with pytest.raises(RuntimeError, match="driver rejected"): + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "driver rejected" in (ledger.statements[0].error or "") + + def test_capture_failure_does_not_break_the_statement(self) -> None: + class _Unprintable: + def __str__(self) -> str: + raise ValueError("cannot render") + + class _Driver: + def __init__(self) -> None: + self.ran = False + + def sql(self, statement): + self.ran = True + + driver = _Driver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql(_Unprintable()) + + assert driver.ran + assert ledger.is_empty() + + +class TestScriptSplitting: + def test_transient_prefix_does_not_launder_a_tuning_statement(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)") + + assert [s.statement for s in ledger.statements] == [ + "SET x=1;", + "CREATE TABLE t (a INT, z INT) ORDER BY (z)", + ] + receipt = corroborate(ledger, None) + verdicts = {entry.statement: entry.verdict for entry in receipt.entries} + assert verdicts["SET x=1;"] == "transient" + assert verdicts["CREATE TABLE t (a INT, z INT) ORDER BY (z)"] != "transient" + + def test_execute_splits_multi_statement_string_under_the_filter(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ()) + + proxy.execute("SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)") + + assert driver.calls == ["SET x=1; CREATE TABLE t (a INT, z INT) ORDER BY (z)"] + assert [s.statement for s in ledger.statements] == ["CREATE TABLE t (a INT, z INT) ORDER BY (z)"] + + def test_readback_prefix_does_not_hide_a_following_mutation(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("SELECT 1; CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [s.statement for s in ledger.statements] == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + + def test_single_statement_is_recorded_unchanged(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + statement = " CREATE TABLE t (a INT) \n" + + proxy.sql(statement) + + assert [s.statement for s in ledger.statements] == [statement] + + +_TRAILING_SQL_COMMENT = "-- trailing comment" +_READBACK_SQL_COMMENT = "-- readback with a trailing comment" + + +class TestCommentOnlyFragments: + def test_trailing_comment_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert driver.calls == [f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}"] + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_trailing_comment_after_verb_call_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + receipt = corroborate(ledger, None) + assert [entry.statement for entry in receipt.entries] == ["CREATE INDEX idx ON t(x);"] + + def test_trailing_comment_after_executescript_is_not_recorded(self) -> None: + class _ScriptDriver: + def __init__(self) -> None: + self.scripts: list[str] = [] + + def executescript(self, script, *args, **kwargs): + self.scripts.append(str(script)) + + driver = _ScriptDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.executescript(f"CREATE INDEX idx ON t(x); {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_readback_script_with_comment_tail_records_nothing(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"SELECT 1; {_READBACK_SQL_COMMENT}") + + assert driver.calls == [f"SELECT 1; {_READBACK_SQL_COMMENT}"] + assert ledger.is_empty() + + def test_empty_split_fragment_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute("CREATE INDEX idx ON t(x);;") + + assert driver.calls == ["CREATE INDEX idx ON t(x);;"] + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_empty_split_fragment_after_verb_call_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL, execute_verbs=("sql",)) + + proxy.sql("CREATE INDEX idx ON t(x);;") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_empty_split_fragment_after_executescript_is_not_recorded(self) -> None: + class _ScriptDriver: + def __init__(self) -> None: + self.scripts: list[str] = [] + + def executescript(self, script, *args, **kwargs): + self.scripts.append(str(script)) + + driver = _ScriptDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.executescript("CREATE INDEX idx ON t(x);;") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + def test_delimiter_and_comment_tail_after_execute_is_not_recorded(self) -> None: + driver = _SqlDriver() + ledger = AppliedTuningLedger() + proxy = recording_connection(driver, ledger, PHASE_DDL) + + proxy.execute(f"CREATE INDEX idx ON t(x);; {_TRAILING_SQL_COMMENT}") + + assert [s.statement for s in ledger.statements] == ["CREATE INDEX idx ON t(x);"] + + +class TestQueryJobs: + def test_submission_alone_is_not_recorded(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts) CLUSTER BY a") + + assert client.submitted + assert ledger.is_empty() + + def test_successful_job_is_recorded_when_it_finishes(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts) CLUSTER BY a") + assert job.job_id == "job-1" + job.result() + job.result() + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + + def test_job_failing_after_submission_is_recorded_failed(self) -> None: + client = _QueryClient(error=RuntimeError("Quota exceeded")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) PARTITION BY DATE(ts)") + with pytest.raises(RuntimeError, match="Quota exceeded"): + job.result() + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "Quota exceeded" in (ledger.statements[0].error or "") + + def test_exception_accessor_settles_the_final_state(self) -> None: + client = _QueryClient(error=RuntimeError("Not found")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + assert str(job.exception()) == "Not found" + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + + def test_timeout_while_waiting_is_not_a_final_state(self) -> None: + from concurrent.futures import TimeoutError as FutureTimeoutError + + client = _QueryClient(error=FutureTimeoutError("still running")) + proxy, ledger = _ddl_proxy(client, ("query",)) + + job = proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + with pytest.raises(FutureTimeoutError): + job.result(timeout=1) + + assert ledger.is_empty() + + def test_dry_run_and_readback_jobs_are_not_recorded(self) -> None: + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, ("query",)) + + dry = proxy.query( + "CREATE TABLE d.t (a INT64) CLUSTER BY a", + job_config=SimpleNamespace(dry_run=True), + ) + select = proxy.query("SELECT * FROM d.t") + dry.result() + select.result() + + assert ledger.is_empty() + + def test_submission_error_is_recorded_failed(self) -> None: + class _Rejecting(_QueryClient): + def query(self, statement, job_config=None, **kwargs): + raise RuntimeError("bad request") + + proxy, ledger = _ddl_proxy(_Rejecting(), ("query",)) + + with pytest.raises(RuntimeError, match="bad request"): + proxy.query("CREATE TABLE d.t (a INT64) CLUSTER BY a") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + + +class TestLazyFrames: + def test_submission_alone_records_nothing(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert driver.calls == ["CREATE TABLE t (a INT) PARTITIONED BY (a)"] + assert ledger.is_empty() + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "noop" + + def test_action_records_executed_once(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + assert ledger.is_empty() + frame.collect() + frame.show() + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + assert "PARTITIONED BY (a)" in ledger.statements[0].statement + + def test_failing_action_records_failed(self) -> None: + driver = _LazySqlDriver(error=RuntimeError("action failed")) + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + assert ledger.is_empty() + with pytest.raises(RuntimeError, match="action failed"): + frame.collect() + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED] + assert "action failed" in (ledger.statements[0].error or "") + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "failed" + + def test_readback_frame_action_is_not_recorded(self) -> None: + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + frame = proxy.sql("SELECT COUNT(*) FROM t") + frame.collect() + + assert len(driver.calls) == 1 + assert ledger.is_empty() + + def test_eager_none_result_still_records_at_submission(self) -> None: + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, ("sql",)) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] + + +class TestExecutescriptFailures: + def test_failing_script_records_failed_per_statement(self) -> None: + class _FailingScriptDriver: + def executescript(self, script, *args, **kwargs): + raise RuntimeError("script boom") + + proxy, ledger = _ddl_proxy(_FailingScriptDriver(), ()) + + with pytest.raises(RuntimeError, match="script boom"): + proxy.executescript("CREATE INDEX i ON t(x); CREATE INDEX j ON t(y);") + + assert [s.status for s in ledger.statements] == [STATEMENT_FAILED, STATEMENT_FAILED] + assert all("script boom" in (s.error or "") for s in ledger.statements) + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "failed" + + +class TestFootprint: + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (a INT) USING parquet PARTITIONED BY (a)", + "CREATE TABLE t (a INT) USING parquet CLUSTERED BY (a) INTO 4 BUCKETS", + "CREATE TABLE t (a INT) DISTKEY(a)", + "CREATE TABLE t (a INT) DISTSTYLE KEY", + "CREATE TABLE t (a INT) DUPLICATE KEY (a) DISTRIBUTED BY HASH (a)", + "CREATE TABLE t (a INT) PRIMARY INDEX (a)", + "CREATE TABLE t (a INT) WITH (distribution = 'hash')", + "CREATE TABLE t (a INT) WITH (partitioning = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (sorted_by = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (bucketed_by = ARRAY['a'], bucket_count = 4)", + ], + ) + def test_new_footprint_keywords_are_schema_tuning(self, statement: str) -> None: + assert is_schema_tuning_statement(statement) + + def test_plain_table_is_still_not_schema_tuning(self) -> None: + assert not is_schema_tuning_statement("CREATE TABLE t (a INT, b INT)") + + def test_keyword_inside_a_literal_does_not_count(self) -> None: + assert not is_schema_tuning_statement("CREATE TABLE t (a INT) COMMENT 'partitioned by day'") + + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (distkey INT)", + "CREATE TABLE t (a INT, distkey TEXT)", + "CREATE TABLE t (diststyle VARCHAR(16))", + "CREATE TABLE t (sortkey INT)", + "CREATE TABLE t (a INT, sorted_by DOUBLE)", + "CREATE TABLE t (bucketed_by INT)", + "CREATE TABLE t (unique INT)", + "CREATE TABLE t (check INT)", + ], + ) + def test_bare_keyword_as_a_column_name_is_not_schema_tuning(self, statement: str) -> None: + assert not is_schema_tuning_statement(statement) + + @pytest.mark.parametrize( + "statement", + [ + "CREATE TABLE t (a INT) PARTITIONED BY (a)", + "CREATE TABLE t (a INT) CLUSTERED BY (a) INTO 4 BUCKETS", + "CREATE TABLE t (a INT) DISTSTYLE EVEN DISTKEY (a)", + "CREATE TABLE t (a INT) DUPLICATE KEY (a)", + "CREATE TABLE t (a INT) PRIMARY INDEX (a)", + "CREATE TABLE t (a INT) WITH (distribution = 'hash')", + "CREATE TABLE t (a INT) WITH (partitioning = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (sorted_by = ARRAY['a'])", + "CREATE TABLE t (a INT) WITH (bucketed_by = ARRAY['a'])", + "CREATE TABLE t (a INT UNIQUE)", + "CREATE TABLE t (a INT, UNIQUE NULLS NOT DISTINCT (a))", + "CREATE TABLE t (a INT, UNIQUE NULLS DISTINCT (a))", + ], + ) + def test_bare_keyword_in_clause_syntax_is_schema_tuning(self, statement: str) -> None: + assert is_schema_tuning_statement(statement) + + +class TestAdapterCreateSchema: + def test_spark_tuned_schema_ddl_reaches_the_ledger(self) -> None: + from benchbox.platforms.spark import SparkAdapter + + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, SparkAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE orders (o_id INT, o_date DATE) USING parquet PARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT) USING delta CLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="parquet", + logger=logging.getLogger("test"), + _remove_orphaned_table_location=lambda *a, **k: None, + ) + + SparkAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(driver.calls) == 3 + recorded = [s.statement for s in ledger.statements] + assert len(recorded) == 2 + assert any("PARTITIONED BY" in s for s in recorded) + assert any("CLUSTER BY" in s for s in recorded) + assert {s.phase for s in ledger.statements} == {PHASE_DDL} + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + + def test_spark_tuned_schema_ddl_with_lazy_frames_reaches_the_ledger(self) -> None: + from benchbox.platforms.spark import SparkAdapter + + driver = _LazySqlDriver() + proxy, ledger = _ddl_proxy(driver, SparkAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE orders (o_id INT, o_date DATE) USING parquet PARTITIONED BY (o_date);\n" + "CREATE TABLE lineitem (l_id INT, l_key INT) USING delta CLUSTER BY (l_key);\n" + "CREATE TABLE plain (p_id INT);" + ), + table_format="parquet", + logger=logging.getLogger("test"), + _remove_orphaned_table_location=lambda *a, **k: None, + ) + + SparkAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(driver.calls) == 3 + assert all(frame.actions for frame in driver.frames) + recorded = [s.statement for s in ledger.statements] + assert len(recorded) == 2 + assert any("PARTITIONED BY" in s for s in recorded) + assert any("CLUSTER BY" in s for s in recorded) + assert {s.phase for s in ledger.statements} == {PHASE_DDL} + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + + def test_bigquery_tuned_schema_ddl_records_final_job_state(self) -> None: + from benchbox.platforms.bigquery import BigQueryAdapter + + client = _QueryClient() + proxy, ledger = _ddl_proxy(client, BigQueryAdapter.ledger_execute_verbs) + fake_self = SimpleNamespace( + dataset_id="ds", + location="US", + _create_schema_with_tuning=lambda *a, **k: ( + "CREATE TABLE t1 (a INT64, ts TIMESTAMP) PARTITION BY DATE(ts) CLUSTER BY a;\nCREATE TABLE t2 (b INT64);" + ), + _convert_to_bigquery_table=lambda statement: statement, + logger=logging.getLogger("test"), + ) + + with patch("benchbox.platforms.bigquery.bigquery"): + BigQueryAdapter.create_schema(fake_self, SimpleNamespace(), proxy) + + assert len(client.submitted) == 2 + assert [s.status for s in ledger.statements] == [EXECUTED] + assert "CLUSTER BY a" in ledger.statements[0].statement + + def test_datafusion_tuned_ddl_through_sql_is_recorded(self) -> None: + from benchbox.platforms.datafusion import DataFusionAdapter + + driver = _SqlDriver() + proxy, ledger = _ddl_proxy(driver, DataFusionAdapter.ledger_execute_verbs) + + proxy.sql("CREATE TABLE t (a INT) PARTITIONED BY (a)") + proxy.sql("SELECT COUNT(*) FROM t") + + assert len(driver.calls) == 2 + assert [(s.phase, s.status) for s in ledger.statements] == [(PHASE_DDL, EXECUTED)] diff --git a/tests/unit/platforms/test_local_platform_ctas_sort.py b/tests/unit/platforms/test_local_platform_ctas_sort.py index 01f7c8e956..133fedde64 100644 --- a/tests/unit/platforms/test_local_platform_ctas_sort.py +++ b/tests/unit/platforms/test_local_platform_ctas_sort.py @@ -93,3 +93,43 @@ def test_datafusion_apply_ctas_sort_orders_rows(self, tmp_path): rows = connection.sql("SELECT l_shipdate, l_orderkey FROM lineitem").to_pydict() assert rows["l_shipdate"] == ["2024-01-01", "2024-02-01", "2024-03-01"] assert rows["l_orderkey"] == [1, 2, 3] + + def test_datafusion_real_ctas_sort_is_recorded_in_ledger(self, tmp_path): + from benchbox.core.tuning.applied_ledger import AppliedTuningLedger + from benchbox.core.tuning.interface import TuningColumn, TuningType + + adapter = DataFusionAdapter(working_dir=str(tmp_path / "df_workdir"), data_format="csv") + adapter._applied_tuning_ledger = AppliedTuningLedger() + connection = adapter.create_connection() + + connection.execute( + "CREATE TABLE lineitem AS " + "SELECT * FROM (VALUES (3, '2024-03-01'), (1, '2024-01-01'), (2, '2024-02-01')) " + "AS t(l_orderkey, l_shipdate)" + ) + + sort_columns = [ + TuningColumn(name="l_shipdate", type="VARCHAR", order=1), + TuningColumn(name="l_orderkey", type="INTEGER", order=2), + ] + table_tuning = Mock() + table_tuning.get_columns_by_type.side_effect = lambda tuning_type: ( + sort_columns if tuning_type == TuningType.SORTING else [] + ) + tuning_config = Mock() + tuning_config.table_tunings = {"lineitem": table_tuning} + + ledger = adapter._applied_tuning_ledger + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "noop" + + assert adapter.apply_ctas_sort("lineitem", tuning_config, connection) is True + + executed = [s for s in ledger.statements if s.status == "executed"] + assert len(executed) == 1 + assert executed[0].phase == "post_load" + assert "ORDER BY l_shipdate, l_orderkey" in executed[0].statement + assert ledger.overall_status(tuning_enabled=True, has_config=True) == "applied_unverified" + + rows = connection.sql("SELECT l_shipdate, l_orderkey FROM lineitem").to_pydict() + assert rows["l_shipdate"] == ["2024-01-01", "2024-02-01", "2024-03-01"] + assert rows["l_orderkey"] == [1, 2, 3] diff --git a/tests/unit/platforms/test_spark_helpers.py b/tests/unit/platforms/test_spark_helpers.py index 8301eaa047..761f1f4b75 100644 --- a/tests/unit/platforms/test_spark_helpers.py +++ b/tests/unit/platforms/test_spark_helpers.py @@ -150,6 +150,53 @@ def test_strips_existing_using_clause(self) -> None: assert "DELTA" not in result.upper() assert result.endswith("USING PARQUET") + @pytest.mark.parametrize("clause", ["PARTITIONED BY (id)", "CLUSTER BY (id)"]) + def test_using_precedes_trailing_tuning_clause(self, clause: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE t (id INT, amount DECIMAL(10,2))\n{clause};", + table_format="delta", + ) + assert result == f"CREATE TABLE t (id INT, amount DECIMAL(10,2)) USING DELTA\n{clause}" + + @pytest.mark.parametrize( + ("columns", "name"), + [ + ("id INT, note STRING DEFAULT ' ) '", "t"), + ("id INT, note STRING DEFAULT 'it''s ) here'", "t"), + ("id INT, note STRING DEFAULT 'a\\' ) b'", "t"), + ("id INT, `we ) ird` STRING", "t"), + ("id INT, amount DECIMAL(10,2), tags MAP>", "`my table`"), + ], + ) + def test_using_stays_outside_column_list_with_quoted_parentheses(self, columns: str, name: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE {name} ({columns}) PARTITIONED BY (id);", + table_format="delta", + ) + assert f"({columns})" in result + assert result.count("USING DELTA") == 1 + assert result.index(f"({columns})") + len(f"({columns})") < result.index("USING DELTA") + assert result.index("USING DELTA") < result.index("PARTITIONED BY") + + @pytest.mark.parametrize( + "columns", + [ + "id INT /* ) */, amount INT", + "id INT /* ( unbalanced */, amount INT", + "id INT -- )\n, amount INT", + "id INT -- ( unbalanced\n, amount INT", + ], + ) + def test_using_stays_outside_column_list_with_comment_parentheses(self, columns: str) -> None: + result = optimize_spark_table_definition( + f"CREATE TABLE t ({columns}) PARTITIONED BY (id);", + table_format="delta", + ) + assert f"({columns})" in result + assert result.count("USING DELTA") == 1 + assert result.index(f"({columns})") + len(f"({columns})") < result.index("USING DELTA") + assert result.index("USING DELTA") < result.index("PARTITIONED BY") + def test_non_create_table_unchanged(self) -> None: sql = "INSERT INTO t VALUES (1)" assert optimize_spark_table_definition(sql, table_format="parquet") == sql