fabric-toolbox/tools/DAXPerformanceTesting/DAXPerformanceTesting.ipynb

DAX Performance Testing

Summary

This notebook is designed to measure DAX query timings under different cache states (cold, warm, and hot).

Requirements:

1. DAX Queries from Excel or YAML

   • You must provide an Excel file containing the DAX queries in a table you wish to test.
   • For each query, a column needs align with the runQueryType used for a given queryId.
   • This notebook reads those queries and executes them on one or more Power BI/Fabric models.
   • For Yaml files please refer to the sample-yaml-dax-query-file.yaml file in the media folder

2. Lakehouse Logging

   • You also must attach the appropriate Lakehouse in Fabric so that logs can be saved (both in a table and as files if you choose).

3. Capacity Pause/Resume

   • In some scenarios (e.g., simulating a "cold" cache on DirectQuery or Import models), the code pauses and resumes capacities.
   • Warning: Pausing a capacity will interrupt any running workloads on that capacity. Resuming will take time and resources, and can affect other workspaces assigned to the same capacity.

Install the latest Semantic Link Labs package

Check here to see the latest version.

In [ ]:

%pip install semantic-link-labs

Import the library and necessary packages

In [ ]:

# Standard Library Imports
import time
import random
import functools
import builtins
import xml.etree.ElementTree as ET
from threading import local
from typing import Any, Callable, Generator, Optional, Type
from contextlib import contextmanager
from uuid import uuid4
from datetime import datetime

# Third-party Imports
import pandas as pd
from pyspark.sql.functions import col, sum as _sum, when, countDistinct

# Local Application/Library-specific Imports
import sempy.fabric as fabric
import sempy_labs as labs
import yaml

Global configurations & variables

In [ ]:

# Define models and their configurations for testing
models = [
    {
        "name": "Model Name", # The name of the semantic model
        "storageMode": "DirectLake",  # Import, DirectQuery, or DirectLake
        "cache_types": ["cold", "warm", "hot"], # List of cache types to be run (hot, warm, and cold)
        "model_workspace_name": "Model Workspace Name", # The workspace name of the semantic model
        "database_name": "Lakehouse Name",  # Only needed for cold cache queries for Import and DirectQuery
        "database_workspace_name": "Lakehouse Workspace Name",  # Only needed for cold cache queries for Import and DirectQuery
        "runQueryType": "query", # The name of the column in your DAX Excel file contains the query to be run
    },
]

# Define DAX queries inline (set to None to load from file instead)
# Each dict must have 'queryId' and one or more query columns matching model 'runQueryType'
dax_queries_inline = None
# dax_queries_inline = [
#     {"queryId": 1, "query": """EVALUATE ROW("Measure", 'Sales'[Amount])"""}, 
#     {"queryId": 2, "query": """EVALUATE ROW("Measure", 'Sales'[Quantity])"""}, 
# ]

# File-based query settings (used when dax_queries_inline is None)
query_file_path = "Files/DAXQueries.xlsx"  # Path to the query file relative to the mount
query_file_mount_path = "/default"              # Mount location where the file is stored
query_worksheet_name = "DAXQueries"          # Worksheet name (for Excel files)

# Only needed for cold cache queries for Import and DirectQuery
workspace_capacities = {
    "Workspace Name": {
        "capacity_name": "Testing Capacity Name",
        "alt_capacity_name": "Alternate Capacity Name",
    }
}

# Additional arguments controlling the behavior of query execution and logging
additional_arguments = {
    "roundNumber": 1, # The current round of DAX testing. Will be considered when determine if maxNumberPerQuery is met or not
    "onlyRunNewQueries": True, # Will determine if queries will stop being tested after maxNumberPerQuery is met
    "maxNumberPerQuery": 1, # The max number of queries to capture per round, queryId, model and cache type
    "maxFailuresBeforeSkipping": 5, # The number of failed query attempts per round, queryId, model and cache type before skipping
    "numberOfRunsPerQueryId": 15, # The number of times to loop over each queryId. If all combos have met maxNumberPerQuery, the loop will break
    "stopQueryIdsAt": 99, # Allows you to stop the queryId loop at a certain number, even if there are more queries present, i.e., there are queryIds 1-20 but stop at 5
    "forceStartQueriesAt1": False, # If set to False, testing will stop at the first incomplete queryId instead of starting at queryId 1  
    "logTableName": "DAXTestingLogTableName", # The name of the table in the attached lakehouse to save the performance logs to
    "clearAllLogs": False, # Will drop the existing logs table before starting testing
    "clearCurrentRoundLogs": False, # Will delete the logs associated with the current roundNumber before starting testing
    "randomizeRuns": True, # Will randomize the model and cache type combos when testing
    "skipSettingHotCache": False, # Should be False if randomizing the runs. If the runs are randomized, the previous warm cache run will set the hot cache
    "pauseAfterSettingCache": 5, # The number of seconds to wait after setting the cache
    "pauseAfterRunningQuery": 5, # The number of second to wait before writing the logs to the log table
    "pauseBetweenRuns": 30, # The number of seconds to wait before starting the next query
}

# Define the expected schema for DAX trace log events
event_schema = {
    "DirectQueryBegin": [
        "EventClass", "CurrentTime", "TextData", "StartTime", 
        "EndTime", "Duration", "CpuTime", "Success", "SessionID"
    ],
    "DirectQueryEnd": [
        "EventClass", "CurrentTime", "TextData", "StartTime", 
        "EndTime", "Duration", "CpuTime", "Success", "SessionID"
    ],
    "VertiPaqSEQueryBegin": [
        "EventClass", "EventSubclass", "CurrentTime", 
        "TextData", "StartTime", "SessionID"
    ],
    "VertiPaqSEQueryEnd": [
        "EventClass", "EventSubclass", "CurrentTime", "TextData", 
        "StartTime", "EndTime", "Duration", "CpuTime", "Success", "SessionID"
    ],
    "VertiPaqSEQueryCacheMatch": [
        "EventClass", "EventSubclass", "CurrentTime", "TextData", "SessionID"
    ],
    "QueryBegin": [
        "EventClass", "EventSubclass", "CurrentTime", "TextData", 
        "StartTime", "ConnectionID", "SessionID", "RequestProperties"
    ],
    "QueryEnd": [
        "EventClass", "EventSubclass", "CurrentTime", "TextData", 
        "StartTime", "EndTime", "Duration", "CpuTime", "Success", 
        "ConnectionID", "SessionID"
    ],
}

# Dictionary to track if a capacity pause is needed for each model during testing
model_pause_capacity_needed = {}

# Variables for Pausing/Resuming Capacities: credentials and configuration parameters for Azure Key Vault and resource management
resource_group_name = ""
subscription_id = ""
key_vault_uri_secret_name = ""
key_vault_client_id_secret_name = ""
key_vault_tenant_id_secret_name = ""
key_vault_client_secret_secret_name = ""

# Enforce case-sensitivity in Spark to ensure column name matching is exact
spark.conf.set("spark.sql.caseSensitive", True)

Logging & Retry Decorators, Basic Helpers

In [ ]:

# Thread-local storage variable to manage indentation for logging.
_thread_local = local()

@contextmanager
def dynamic_indented_print() -> Generator[None, None, None]:
    """
    A context manager that overrides the built-in print function to automatically
    indent messages based on the current call depth. This helps visualize nested logs.
    """
    original_print = builtins.print

    def custom_print(*args: Any, **kwargs: Any) -> None:
        # Use current call depth (default 0) to set indentation.
        depth = getattr(_thread_local, "call_depth", 0)
        indent = "    " * depth
        original_print(indent + " ".join(map(str, args)), **kwargs)

    builtins.print = custom_print
    try:
        yield
    finally:
        builtins.print = original_print

def log_function_calls(func: Callable) -> Callable:
    """
    Decorator that logs the start and end of a function call.
    It increases the indentation for nested function calls for better readability.
    """
    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        if not hasattr(_thread_local, "call_depth"):
            _thread_local.call_depth = 0

        with dynamic_indented_print():
            print(f"✅ {func.__name__} - Starting")
            _thread_local.call_depth += 1  # Increase indentation level for inner calls
            try:
                result = func(*args, **kwargs)
            finally:
                _thread_local.call_depth -= 1  # Decrease indentation level when done
                print(f"✅ {func.__name__} - Ending")
        return result

    return wrapper

def retry(exceptions: tuple[Type[Exception], ...],
          tries: int = 3,
          delay: int = 5,
          backoff: int = 2,
          logger: Callable = print) -> Callable:
    """
    A decorator for retrying a function call with exponential backoff.
    This is useful for transient errors (e.g., waiting for a capacity status update).
    """
    def decorator_retry(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper_retry(*args: Any, **kwargs: Any) -> Any:
            _tries, _delay = tries, delay
            while _tries > 1:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    with dynamic_indented_print():
                        print(f"⚠️ {func.__name__} failed with {e}, retrying in {_delay} seconds...")
                    time.sleep(_delay)
                    _tries -= 1
                    _delay *= backoff  # Increase the delay for each retry
            return func(*args, **kwargs)
        return wrapper_retry
    return decorator_retry

Helper Functions

In [ ]:

def to_camel_case(s: str) -> str:
    """Converts a string to camel case (e.g., 'Event Class' -> 'eventClass')."""
    words = s.split()
    if not words:
        return s
    return words[0].lower() + ''.join(word.capitalize() for word in words[1:])

@retry(Exception, tries=10, delay=2, backoff=2, logger=print)
def wait_for_trace_start(trace_connection, trace_name: str) -> bool:
    """Polls until the trace with the specified name is detected as started."""
    traces = trace_connection.list_traces()
    if traces.loc[traces["Name"] == trace_name].shape[0] == 0:
        raise Exception("Trace has not started yet")
    return True

def extract_session_id_from_query_begin(logs: pd.DataFrame) -> Optional[str]:
    """
    Extracts the Session ID from QueryBegin events that include 'SemPy' in their Request Properties.
    Pre-filters with str.contains to avoid unnecessary XML parsing.
    """
    query_begins = logs[
        (logs["Event Class"] == "QueryBegin") &
        (logs["Request Properties"].str.contains("SemPy", na=False))
    ]
    for req_props, session_id in zip(query_begins["Request Properties"], query_begins["Session ID"]):
        try:
            root = ET.fromstring(req_props)
            for child in root:
                if child.tag.endswith("SspropInitAppName") and child.text == "SemPy":
                    return session_id
        except Exception as e:
            print(f"Error parsing Request Properties XML: {e}")
    return None

@retry(Exception, tries=60, delay=3, backoff=1, logger=print)
def wait_for_query_end_event(trace) -> tuple:
    """
    Waits until a QueryBegin event with 'SemPy' is detected and its corresponding QueryEnd event is present.
    Returns a tuple of (session_id, logs).
    """
    logs = trace.get_trace_logs()
    session_id = extract_session_id_from_query_begin(logs)
    if not session_id:
        raise Exception("QueryBegin event with SspropInitAppName 'SemPy' not found yet")

    query_end_events = logs[(logs["Event Class"] == "QueryEnd") & (logs["Session ID"] == session_id)]
    if query_end_events.empty:
        raise Exception("QueryEnd event for Session ID not collected yet")
    return session_id, logs

def collect_filtered_trace_logs(trace) -> pd.DataFrame:
    """Waits for QueryEnd, stops the trace, and returns logs filtered to the relevant Session ID."""
    session_id, _ = wait_for_query_end_event(trace)
    logs = trace.stop()
    return logs[logs["Session ID"] == session_id]

def execute_traced_query(model: dict, expression: str, trace_label: str = "Trace") -> tuple:
    """
    Execute a DAX query with AS trace logging.
    Handles the full trace lifecycle: drop old traces, create new trace, start, execute query, collect logs.
    Returns (result, filtered_logs) where result is "Success" or the error string.
    """
    fabric.create_trace_connection(model["name"], model["model_workspace_name"]).drop_traces()
    trace_name = f"{trace_label} {uuid4()}"
    with fabric.create_trace_connection(model["name"], model["model_workspace_name"]) as trace_conn:
        with trace_conn.create_trace(event_schema, trace_name) as trace:
            print(f"🔍 Starting trace: {trace_label}")
            trace.start()
            wait_for_trace_start(trace_conn, trace_name)
            result = "Success"
            try:
                print("⚡ Executing DAX query")
                fabric.evaluate_dax(model["name"], expression, workspace=model["model_workspace_name"])
                print("✅ DAX query succeeded")
            except Exception as e:
                result = str(e)
                print(f"❌ DAX query failed: {e}")
            print("📜 Collecting trace logs")
            filtered_logs = collect_filtered_trace_logs(trace)
            return result, filtered_logs

@retry(Exception, tries=30, delay=5, backoff=1, logger=print)
def check_model_online(model: dict) -> None:
    """Checks if the given model is online by executing a simple DAX query."""
    dax_query_eval_1(model)

def get_capacity_status(_capacity_name: str) -> str:
    """Retrieves the current status of a capacity by its name."""
    capacity_info = labs.list_capacities()
    return capacity_info.loc[
        capacity_info["Display Name"] == _capacity_name, "State"
    ].iloc[0]

@retry(Exception, tries=30, delay=2, backoff=2, logger=print)
def wait_for_capacity_status(_capacity_name: str, target_status: str) -> str:
    """Polls until the capacity status matches the target status."""
    current_status = get_capacity_status(_capacity_name)
    if current_status != target_status:
        raise Exception("Capacity status not updated yet")
    return current_status

def dax_query_eval_1(model: dict) -> None:
    """Executes a simple DAX query ("EVALUATE {1}") to verify connectivity."""
    try:
        fabric.evaluate_dax(
            model["name"], "EVALUATE {1}", workspace=model["model_workspace_name"]
        )
    except Exception as e:
        raise Exception("Failed to query model") from e

@log_function_calls
def wait_for_model_to_come_online(model: dict) -> None:
    """Waits until the model is online by calling check_model_online."""
    try:
        check_model_online(model)
        print("✅ Model is online")
    except Exception as e:
        raise Exception("❌ Model failed to come online") from e

@log_function_calls
def load_dax_queries(file_path: str, mount_path: str, worksheet_name: Optional[str] = None, inline_queries: Optional[list] = None) -> pd.DataFrame:
    """
    Loads DAX queries from inline list, Excel, or YAML file.
    If inline_queries is provided, uses those directly.
    Otherwise loads from file. The first column must be 'queryId'.
    """
    if inline_queries is not None:
        if not inline_queries:
            raise ValueError("dax_queries_inline is set but empty. Provide at least one query or set to None to load from file.")
        df = pd.DataFrame(inline_queries)
        if "queryId" not in df.columns:
            raise ValueError(f"Inline queries missing required 'queryId' column. Found columns: {list(df.columns)}")
        print(f"📋 Loading {len(inline_queries)} inline DAX queries")
        return df

    file_type = file_path.split('.')[-1].lower()
    full_path = f"{notebookutils.fs.getMountPath(mount_path)}/{file_path}"

    if file_type == 'xlsx':
        return pd.read_excel(full_path, sheet_name=worksheet_name)
    elif file_type in ['yml', 'yaml']:
        with open(full_path, 'r') as f:
            data = yaml.safe_load(f)
        if not data:
            raise ValueError(f"No queries loaded from '{full_path}'. Check the file is not empty.")
        return pd.DataFrame(data)
    else:
        raise ValueError(f"Unsupported file type: {file_type}")

Pause & Resume Capacity

In [ ]:

@log_function_calls
def update_model_pause_status(event: str, model: dict = None, workspace: str = None) -> None:
    """
    Updates the dictionary that tracks if a capacity pause is needed for each model.
    Different events trigger different updates:
      - "initialize": Set default based on storageMode.
      - "model_queried": Mark model as needing a pause for Import/DirectQuery.
      - "capacity_paused": Reset the pause flag after a capacity operation.
    """
    if event == "initialize":
        for m in models:
            model_pause_capacity_needed[m["name"]] = False if m["storageMode"] == "DirectLake" else True

    elif event == "model_queried" and model is not None:
        if model["storageMode"] == "Import":
            model_pause_capacity_needed[model["name"]] = True
        elif model["storageMode"] == "DirectQuery":
            target_db = model["database_name"]
            target_db_workspace = model["database_workspace_name"]
            for m in models:
                if (
                    m["storageMode"] == "DirectQuery" and 
                    m["database_name"] == target_db and 
                    m["database_workspace_name"] == target_db_workspace
                ):
                    model_pause_capacity_needed[m["name"]] = True

    elif event == "capacity_paused" and workspace is not None:
        # Reset pause flags after capacity operations complete
        for m in models:
            if (
                (m["storageMode"] == "Import" and m["model_workspace_name"] == workspace) or
                (m["storageMode"] == "DirectQuery" and m["database_workspace_name"] == workspace)
            ):
                model_pause_capacity_needed[m["name"]] = False
    else:
        print(f"⚠️ Unknown event '{event}' or missing required parameter(s).")

    print("📝 Updated model_pause_capacity_needed")

@log_function_calls
@retry(Exception, tries=10, delay=5, backoff=2, logger=print)
def pause_resume_capacity(_capacity_name, _action):
    """
    Pauses or resumes a capacity based on the _action parameter.
    Uses Semantic Link Labs functions along with key vault authentication.
    """
    print(f"🔄 {_action.title()} capacity '{_capacity_name}': Attempting")

    # Check the current status of the capacity.
    current_status = get_capacity_status(_capacity_name)
    
    # Mapping for actions with expected and target statuses.
    action_options = {
        "pause": {
            "expected_status": "Active",  # Can only pause if capacity is active.
            "target_status": "Suspended",
            "action_function": labs.suspend_fabric_capacity,
        },
        "resume": {
            "expected_status": "Suspended",  # Can only resume if capacity is suspended.
            "target_status": "Active",
            "action_function": labs.resume_fabric_capacity,
        },
    }
    
    if current_status == action_options[_action]["expected_status"]:
        print(f"🛠️ {_action.title()} capacity '{_capacity_name}': Requesting action")
        
        # Authenticate using service principal details from the key vault.
        with labs.service_principal_authentication(
            key_vault_uri=key_vault_uri_secret_name,
            key_vault_tenant_id=key_vault_tenant_id_secret_name,
            key_vault_client_id=key_vault_client_id_secret_name,
            key_vault_client_secret=key_vault_client_secret_secret_name
        ):
            # Execute the pause or resume action.
            action_options[_action]["action_function"](
                capacity_name=_capacity_name,
                azure_subscription_id=subscription_id,
                resource_group=resource_group_name
            )
        
        # Wait for the capacity status to reach the target status.
        try:
            wait_for_capacity_status(_capacity_name, action_options[_action]["target_status"])
            print(f"✅ {_action.title()} capacity '{_capacity_name}': Action successful")
        except Exception as e:
            print(f"⚠️ {_action.title()} capacity '{_capacity_name}': Timeout waiting for target status. Error: {e}")
    else:
        print(f"ℹ️ {_action.title()} capacity '{_capacity_name}': Already '{current_status}'")

Cache-setting functions

In [ ]:

@log_function_calls
@retry(exceptions=(Exception,), tries=5, delay=5, backoff=2)
def clear_vertipaq_cache(model: dict) -> None:
    """Clears the VertiPaq cache and verifies with a simple DAX query."""
    print("🧹 Clearing VertiPaq cache")
    wait_for_model_to_come_online(model)
    try:
        labs.clear_cache(model["name"], workspace=model["model_workspace_name"])
        dax_query_eval_1(model)
        print("✅ Clear VertiPaq cache successful")
    except Exception as e:
        print("🔄 Clearing VertiPaq cache failed; retrying...")
        fabric.refresh_tom_cache(model["model_workspace_name"])
        raise e
    time.sleep(5)

@log_function_calls
def set_hot_cache(model: dict, expression: str, successful_query_count_goal: int = 1) -> bool:
    """
    Primes the hot cache by executing the given query multiple times.
    Returns True if the required number of successful executions is reached.
    """
    print("🔥 Setting Hot Cache")
    successful_query_count = 0
    number_of_query_attempts = successful_query_count_goal * 5 if successful_query_count_goal > 1 else 1

    if additional_arguments["skipSettingHotCache"]:
        successful_query_count = successful_query_count_goal
    else:
        for _ in range(number_of_query_attempts):
            result, _ = execute_traced_query(model, expression, "Cache Trace")
            if result == "Success":
                successful_query_count += 1
            if successful_query_count >= successful_query_count_goal:
                break

    print(f"✅ Hot cache set; goal: {successful_query_count_goal} successful queries")
    return successful_query_count >= successful_query_count_goal

@log_function_calls
def set_warm_cache(model: dict, expression: str) -> bool:
    """Sets a warm cache: prime hot cache then clear VertiPaq. For DirectQuery, sets cold cache first."""
    print("🔥 Setting Warm Cache")
    if model["storageMode"] == "DirectQuery":
        set_cold_cache(model)
    hot_cache_successful = set_hot_cache(model, expression)
    clear_vertipaq_cache(model)
    print("✅ Warm cache set")
    return hot_cache_successful

@retry(exceptions=(Exception,), tries=5, delay=5, backoff=2)
def _refresh_tom_cache(workspace_name: str) -> None:
    """Refreshes the TOM cache for the specified workspace."""
    print(f"⌛ Refreshing TOM cache for workspace '{workspace_name}'")
    fabric.refresh_tom_cache(workspace_name)
    print("✅ TOM cache refreshed successfully")

@retry(exceptions=(Exception,), tries=30, delay=3, backoff=2)
def _wait_for_refresh_to_complete(model: dict, refresh_id: str) -> None:
    """Waits for a dataset refresh to complete or fail by checking its status."""
    status = fabric.get_refresh_execution_details(
        model["name"],
        refresh_id,
        workspace=model["model_workspace_name"],
    ).status

    if status not in ["Completed", "Failed"]:
        raise Exception(f"Refresh status is '{status}'; not done yet.")
    print(f"✅ Refresh status: '{status}' - finishing polling.")

@log_function_calls
def set_cold_cache(model: dict) -> None:
    """
    Sets cold cache by clearing all caching.
    For Import/DirectQuery: pauses/resumes capacities.
    For DirectLake: performs clearValues + full refresh.
    """
    print("❄️ Setting Cold Cache")
    if model["storageMode"] != "DirectLake":
        if model_pause_capacity_needed.get(model["name"], False):
            ws_caps = workspace_capacities[model["model_workspace_name"]]
            print(f"🔄 Assigning alternate capacity for workspace '{model['model_workspace_name']}'")
            labs.assign_workspace_to_capacity(ws_caps["alt_capacity_name"], model["model_workspace_name"])
            print(f"✅ Alternate capacity assigned: {ws_caps['alt_capacity_name']}")

            pause_resume_capacity(ws_caps["capacity_name"], "pause")
            pause_resume_capacity(ws_caps["capacity_name"], "resume")

            print(f"🔄 Reassigning primary capacity for workspace '{model['model_workspace_name']}'")
            labs.assign_workspace_to_capacity(ws_caps["capacity_name"], model["model_workspace_name"])
            print(f"✅ Primary capacity assigned: {ws_caps['capacity_name']}")

            _refresh_tom_cache(model["model_workspace_name"])
            wait_for_model_to_come_online(model)
            update_model_pause_status("capacity_paused", workspace=model["model_workspace_name"])
            time.sleep(30)
            clear_vertipaq_cache(model)
    else:
        print("ℹ️ Performing clear refresh for cold cache")
        refresh_status_clear = fabric.refresh_dataset(
            model["name"],
            refresh_type="clearValues",
            workspace=model["model_workspace_name"],
        )
        _wait_for_refresh_to_complete(model, refresh_status_clear)
        print("✅ Clear refresh completed; performing full refresh")
        refresh_status_full = fabric.refresh_dataset(
            model["name"],
            refresh_type="full",
            workspace=model["model_workspace_name"],
        )
        _wait_for_refresh_to_complete(model, refresh_status_full)
        clear_vertipaq_cache(model)
    print("✅ Cold cache set")

Log-table helpers & query checks

In [ ]:

def get_log_table(table_name: str):
    """Retrieves the log table filtered to the current round and QueryEnd events."""
    try:
        raw_table = spark.table(table_name)
        base_filters = (
            (col("roundNumber") == additional_arguments["roundNumber"]) &
            (col("eventClass") == "QueryEnd")
        )
        return raw_table.filter(base_filters)
    except Exception:
        return None

def _get_completed_counts(log_table) -> dict:
    """Query the delta table once to get success/failure counts per (modelName, cacheType, queryId)."""
    if log_table is None:
        return {}
    rows = log_table.groupBy("modelName", "cacheType", "queryId").agg(
        _sum(when(col("success") == "Success", 1).otherwise(0)).alias("success_count"),
        _sum(when(col("success") == "Failure", 1).otherwise(0)).alias("failure_count")
    ).collect()
    return {
        (row["modelName"], row["cacheType"], int(row["queryId"])): {
            "success": row["success_count"], "failure": row["failure_count"]
        }
        for row in rows
    }


def _is_combo_complete(completed_counts: dict, model_name: str, cache_type: str, query_id: int) -> bool:
    """Check if a (model, cache_type, queryId) combo has met its success/failure threshold."""
    key = (model_name, cache_type, query_id)
    if key not in completed_counts:
        return False
    counts = completed_counts[key]
    return (
        counts["success"] >= additional_arguments["maxNumberPerQuery"]
        or counts["failure"] >= additional_arguments["maxFailuresBeforeSkipping"]
    )
def get_starting_query_id(log_table, additional_args: dict) -> int:
    """Determines the next queryId to start from based on existing logs."""
    print("🔍 Determining starting query ID")
    if log_table is not None:
        success_failure_counts = log_table.groupBy("modelName", "cacheType", "queryId").agg(
            _sum(when(col("success") == "Success", 1).otherwise(0)).alias("success_count"),
            _sum(when(col("success") == "Failure", 1).otherwise(0)).alias("failure_count")
        )

        valid_queries = success_failure_counts.filter(
            (col("success_count") >= additional_args["maxNumberPerQuery"]) |
            (col("failure_count") >= additional_args["maxFailuresBeforeSkipping"])
        )

        distinct_combos_count = valid_queries.select("modelName", "cacheType").distinct().count()

        valid_query_ids = valid_queries.groupBy("queryId").agg(
            countDistinct("modelName", "cacheType").alias("valid_combinations")
        ).filter(col("valid_combinations") == distinct_combos_count)

        query_id_list = sorted([row["queryId"] for row in valid_query_ids.select("queryId").collect()])
        max_query_id = 0
        for i, qid in enumerate(query_id_list):
            if qid != i + 1:
                break
            max_query_id = qid

        starting_query_id = 1 if max_query_id == 0 else max_query_id + 1
    else:
        print(f"ℹ️ Log table {additional_args['logTableName']} does not exist")
        starting_query_id = 1

    print(f"✅ Starting query ID set to {starting_query_id}")
    return starting_query_id

Main DAX testing orchestration functions

In [ ]:

@log_function_calls
def run_dax_query_and_collect_logs(model_cache_combo: dict, dax_query: pd.Series, log_table) -> None:
    """
    Runs a single DAX query based on the provided model, cache type, and queryId.
    Handles cache setup, traced execution, and appends results to the log table.
    Pre-filtering of completed/invalid combos is done by the caller.
    """
    model = model_cache_combo["model"]
    sent_dax_expression = dax_query[model["runQueryType"]]
    query_run_name = f"Model: {model['name']}, QueryId: {dax_query['queryId']}, Cache Type: {model_cache_combo['cache_type']}"

    print(f"🚀 Starting query: {query_run_name}")

    set_cache_start_time = datetime.now().isoformat()
    wait_for_model_to_come_online(model)

    # Set the cache according to the type specified
    if model_cache_combo["cache_type"] == "cold":
        set_cold_cache(model)
        cache_set = True
    elif model_cache_combo["cache_type"] == "warm":
        cache_set = set_warm_cache(model, sent_dax_expression)
    else:
        cache_set = set_hot_cache(model, sent_dax_expression)

    time.sleep(additional_arguments["pauseAfterSettingCache"])
    set_cache_end_time = datetime.now().isoformat()
    update_model_pause_status("model_queried", model=model)
    query_start_time = datetime.now().isoformat()

    # Execute the DAX query with trace logging
    dax_query_result, filtered_logs = execute_traced_query(model, sent_dax_expression, "Simple DAX Trace")

    time.sleep(additional_arguments["pauseAfterRunningQuery"])
    query_end_time = datetime.now().isoformat()

    # Capture request properties from the QueryBegin event
    query_begin_rows = filtered_logs[filtered_logs["Event Class"] == "QueryBegin"]
    request_properties_value = query_begin_rows.iloc[0]["Request Properties"] if not query_begin_rows.empty else None
    filtered_logs["Request Properties"] = request_properties_value

    # Convert all column names to camelCase
    filtered_logs.columns = [to_camel_case(c) for c in filtered_logs.columns]

    # Append metadata columns to the trace logs
    filtered_logs = filtered_logs.assign(
        sentExpression=sent_dax_expression,
        setCacheStartTime=set_cache_start_time,
        setCacheEndTime=set_cache_end_time,
        queryStartTime=query_start_time,
        queryEndTime=query_end_time,
        modelName=model["name"],
        queryId=dax_query["queryId"],
        runQueryType=model["runQueryType"],
        cacheType=model_cache_combo["cache_type"],
        queryResult=dax_query_result,
        storageMode=model["storageMode"],
        roundNumber=additional_arguments["roundNumber"],
    )

    if not cache_set:
        print("❌ Cache was not set properly; marking query as failed")
        filtered_logs = filtered_logs.assign(success="Failure")

    # Convert to Spark DataFrame and append to the log table
    filtered_logs_df = spark.createDataFrame(filtered_logs)

    print("💾 Appending logs to table")
    filtered_logs_df.write.format("delta").mode("append").option("mergeSchema", "true").saveAsTable(additional_arguments["logTableName"])
    print("✅ Logs appended to table")
    print("ℹ️ Pausing between runs")
    time.sleep(additional_arguments["pauseBetweenRuns"])

def _initialize_log_table() -> tuple:
    """Initialize the log table based on additional_arguments settings. Returns (log_table, start_query_id)."""
    if additional_arguments["clearCurrentRoundLogs"]:
        print(f"🗑️ Dropping round {additional_arguments['roundNumber']} logs from {additional_arguments['logTableName']}")
        try:
            spark.sql(f"DELETE FROM {additional_arguments['logTableName']} WHERE roundNumber = {additional_arguments['roundNumber']}")
        except Exception:
            print("⚠️ Log table does not exist yet; nothing to clear")

    if additional_arguments["clearAllLogs"]:
        print(f"🗑️ Dropping entire table {additional_arguments['logTableName']}")
        spark.sql(f"DROP TABLE IF EXISTS {additional_arguments['logTableName']}")
        return None, 1

    print(f"🔍 Retrieving table {additional_arguments['logTableName']}")
    log_table = get_log_table(additional_arguments["logTableName"])
    if additional_arguments["clearCurrentRoundLogs"] or additional_arguments["forceStartQueriesAt1"] or log_table is None:
        start_at = 1
    else:
        start_at = get_starting_query_id(log_table, additional_arguments)
    return log_table, start_at

def _validate_and_resume_capacities() -> bool:
    """Check if capacity pause/resume logic is required and resume capacities if so. Returns whether logic is needed."""
    include_pause_resume = any(
        model["storageMode"] in ["Import", "DirectQuery"] and "cold" in model["cache_types"]
        for model in models
    )
    if include_pause_resume:
        for m in models:
            ws = m["model_workspace_name"]
            if ws not in workspace_capacities:
                raise ValueError(f"The workspace '{ws}' (in model '{m['name']}') is not defined in workspace_capacities.")
        for ws, caps in workspace_capacities.items():
            pause_resume_capacity(caps["capacity_name"], "resume")
            pause_resume_capacity(caps["alt_capacity_name"], "resume")
            print(f"✅ Assigning primary capacity '{caps['capacity_name']}' to workspace '{ws}'")
            labs.assign_workspace_to_capacity(caps["capacity_name"], ws)
    return include_pause_resume

def _build_model_cache_combos() -> list:
    """Build the list of (model, cache_type) combinations, optionally randomized."""
    combos = [
        {"model": model, "cache_type": cache_type}
        for model in models
        for cache_type in ["cold", "warm", "hot"]
    ]
    if additional_arguments["randomizeRuns"]:
        print("🔀 Randomizing run order of (model, cache_type)")
        random.shuffle(combos)
    else:
        cache_order = {"cold": 0, "warm": 1, "hot": 2}
        combos.sort(key=lambda c: (c["model"]["name"], cache_order[c["cache_type"]]))
    return combos

@log_function_calls
def run_dax_queries() -> None:
    """
    Main entry point for running all DAX queries.
    Pre-computes completed combos to minimize delta table reads, then executes only remaining work.
    """
    print("🚀 Starting all DAX queries")

    dax_queries = load_dax_queries(query_file_path, query_file_mount_path, query_worksheet_name, dax_queries_inline)

    # Validate that each model's runQueryType matches a column in the query data
    for m in models:
        if m["runQueryType"] not in dax_queries.columns:
            raise ValueError(f"Model '{m['name']}' has runQueryType '{m['runQueryType']}' but query data only has columns: {list(dax_queries.columns)}")

    update_model_pause_status("initialize")

    log_table, start_query_id = _initialize_log_table()
    include_pause_resume_logic = _validate_and_resume_capacities()

    # Pre-compute completed counts (1 Spark query instead of N per-combo queries)
    completed_counts = _get_completed_counts(log_table) if additional_arguments["onlyRunNewQueries"] else {}

    for _, dax_query in dax_queries.iterrows():
        qid = dax_query["queryId"]
        if not (start_query_id <= qid <= additional_arguments["stopQueryIdsAt"]):
            continue

        for _ in range(additional_arguments["numberOfRunsPerQueryId"]):
            all_combos = _build_model_cache_combos()

            # Pre-filter: remove invalid cache types and already-completed combos
            runnable = [
                c for c in all_combos
                if c["cache_type"] in c["model"]["cache_types"]
                and not _is_combo_complete(completed_counts, c["model"]["name"], c["cache_type"], qid)
            ]

            if not runnable:
                print(f"ℹ️ All combos complete for queryId {qid}; skipping remaining runs")
                break

            print(f"▶️ Running {len(runnable)} of {len(all_combos)} combos for queryId {qid}")

            for combo in runnable:
                if include_pause_resume_logic:
                    update_model_pause_status("model_queried", model=combo["model"])
                try:
                    run_dax_query_and_collect_logs(combo, dax_query, log_table)
                except Exception as e:
                    print(f"\u274c Query failed for {combo['model']['name']}/{combo['cache_type']}: {e}")
                    print("\u26a0\ufe0f Continuing to next combo")

            # Refresh log table and completed counts after each round
            print(f"🔄 Refreshing log table from {additional_arguments['logTableName']}")
            log_table = get_log_table(additional_arguments["logTableName"])
            if additional_arguments["onlyRunNewQueries"]:
                completed_counts = _get_completed_counts(log_table)

    if include_pause_resume_logic:
        for ws, caps in workspace_capacities.items():
            pause_resume_capacity(caps["alt_capacity_name"], "pause")
    print("✅ All queries complete")

Execute main flow

In [ ]:

run_dax_queries()

Stop session

In [ ]:

mssparkutils.session.stop()