diff --git a/src/sempy_labs/__init__.py b/src/sempy_labs/__init__.py index 2b2086c00..8ce6bd746 100644 --- a/src/sempy_labs/__init__.py +++ b/src/sempy_labs/__init__.py @@ -234,6 +234,7 @@ evaluate_dax_impersonation, get_dax_query_dependencies, get_dax_query_memory_size, + capture_dax_query_timings, # dax_perf_test, # dax_perf_test_bulk, ) @@ -519,6 +520,7 @@ "list_managed_private_endpoint_fqdns", "get_dax_query_dependencies", "get_dax_query_memory_size", + "capture_dax_query_timings", "get_mirrored_database_definition", "get_mirroring_status", "list_mirrored_databases", diff --git a/src/sempy_labs/_dax.py b/src/sempy_labs/_dax.py index 97ea06a0f..35373e3b0 100644 --- a/src/sempy_labs/_dax.py +++ b/src/sempy_labs/_dax.py @@ -10,7 +10,7 @@ from sempy_labs._model_dependencies import get_model_calc_dependencies from typing import Optional, List, Tuple from sempy._utils._log import log -from uuid import UUID +from uuid import UUID, uuid4 from sempy_labs.directlake._warm_cache import _put_columns_into_memory import sempy_labs._icons as icons import time @@ -285,6 +285,291 @@ def get_dax_query_memory_size( return df["Total Size"].sum() +@log +def capture_dax_query_timings( + dataset: str | UUID, + dax_query: str, + workspace: Optional[str | UUID] = None, + clear_cache: bool = True, + post_wait_seconds: float = 1.2, +) -> pd.DataFrame: + """ + Runs a DAX query with an SSAS trace and returns a one-row summary dataframe. + + Parameters + ---------- + dataset : str | uuid.UUID + Name or ID of the semantic model. + dax_query : str + The DAX query to execute. + workspace : str | uuid.UUID, default=None + The Fabric workspace name or ID. + Defaults to None which resolves to the workspace of the attached lakehouse + or if no lakehouse attached, resolves to the workspace of the notebook. + clear_cache : bool, default=True + If True, clears semantic model cache before query execution. + post_wait_seconds : float, default=1.2 + Wait time after query execution before stopping trace. + + Returns + ------- + pandas.DataFrame + A single-row summary containing total, storage engine, and formula engine timings. + """ + + from sempy_labs._clear_cache import clear_cache as _clear_cache + + event_schema = { + "QueryBegin": [ + "EventClass", + "CurrentTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "QueryEnd": [ + "EventClass", + "CurrentTime", + "Duration", + "CpuTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "VertiPaqSEQueryBegin": [ + "EventClass", + "CurrentTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "VertiPaqSEQueryEnd": [ + "EventClass", + "CurrentTime", + "Duration", + "CpuTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "VertiPaqSEQueryCacheMatch": [ + "EventClass", + "CurrentTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "DirectQueryBegin": [ + "EventClass", + "CurrentTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + "DirectQueryEnd": [ + "EventClass", + "CurrentTime", + "Duration", + "CpuTime", + "SessionID", + "ActivityID", + "RequestID", + "TextData", + ], + } + + if clear_cache: + _clear_cache(dataset=dataset, workspace=workspace) + + run_id = uuid4().hex + run_marker = f"-- RUNID:{run_id}" + tagged_query = f"{run_marker}\n{dax_query}" + + t0 = time.perf_counter() + eval_t1 = t0 + with fabric.create_trace_connection(dataset=dataset, workspace=workspace) as trace_connection: + with trace_connection.create_trace(event_schema) as trace: + trace.start() + _ = fabric.evaluate_dax( + dataset=dataset, + workspace=workspace, + dax_string=tagged_query, + ) + eval_t1 = time.perf_counter() + time.sleep(post_wait_seconds) + df = trace.stop() + + if df is None or df.empty: + elapsed = round((eval_t1 - t0) * 1000.0, 3) + summary = { + "Dataset": dataset, + "Workspace": workspace, + "Total Elapsed (ms)": elapsed, + "Storage Engine (ms)": 0.0, + "Formula Engine (ms)": elapsed, + "Trace Event Count": 0, + "Note": "No trace events returned.", + } + return pd.DataFrame([summary]) + + event_class_col = None + for c in ["Event Class", "EventClass"]: + if c in df.columns: + event_class_col = c + break + + text_col = None + for c in ["Text Data", "TextData"]: + if c in df.columns: + text_col = c + break + + duration_col = "Duration" if "Duration" in df.columns else None + + current_time_col = None + for c in ["Current Time", "CurrentTime"]: + if c in df.columns: + current_time_col = c + break + + session_col = None + for c in ["Session ID", "SessionID"]: + if c in df.columns: + session_col = c + break + + activity_col = None + for c in ["Activity ID", "ActivityID"]: + if c in df.columns: + activity_col = c + break + + request_col = None + for c in ["Request ID", "RequestID"]: + if c in df.columns: + request_col = c + break + + if event_class_col is None: + raise ValueError( + f"{icons.red_dot} Trace output does not contain event class column. " + f"Columns: {list(df.columns)}" + ) + + if current_time_col is not None: + df[current_time_col] = pd.to_datetime(df[current_time_col], errors="coerce") + + query_end_mask = df[event_class_col].astype(str).eq("QueryEnd") + if text_col is not None: + query_end_mask = query_end_mask & df[text_col].astype(str).str.contains( + run_marker, regex=False, na=False + ) + + target = df.loc[query_end_mask].tail(1) + if target.empty: + target = df.loc[df[event_class_col].astype(str).eq("QueryEnd")].tail(1) + + corr = pd.Series(True, index=df.index) + if not target.empty: + target_row = target.iloc[0] + + target_session = ( + str(target_row[session_col]) + if session_col and pd.notna(target_row[session_col]) + else "" + ) + target_request = ( + str(target_row[request_col]) + if request_col and pd.notna(target_row[request_col]) + else "" + ) + target_activity = ( + str(target_row[activity_col]) + if activity_col and pd.notna(target_row[activity_col]) + else "" + ) + + if target_request and request_col is not None: + corr = corr & (df[request_col].astype(str).eq(target_request)) + if session_col is not None and target_session: + corr = corr & (df[session_col].astype(str).eq(target_session)) + elif target_activity and activity_col is not None: + corr = corr & (df[activity_col].astype(str).eq(target_activity)) + if session_col is not None and target_session: + corr = corr & (df[session_col].astype(str).eq(target_session)) + elif target_session and session_col is not None: + corr = corr & (df[session_col].astype(str).eq(target_session)) + + trace_df = df.loc[corr].copy() + + if duration_col is not None: + trace_df[duration_col] = pd.to_numeric(trace_df[duration_col], errors="coerce") + + total_ms = round((eval_t1 - t0) * 1000.0, 3) + if not target.empty and duration_col is not None: + qe_duration = pd.to_numeric(target[duration_col], errors="coerce").iloc[-1] + if pd.notna(qe_duration) and float(qe_duration) > 0: + total_ms = round(float(qe_duration), 3) + + se_end_classes = {"VertiPaqSEQueryEnd", "DirectQueryEnd"} + se_df = trace_df.loc[ + trace_df[event_class_col].astype(str).isin(se_end_classes) + ].copy() + + intervals: List[Tuple[pd.Timestamp, pd.Timestamp]] = [] + if current_time_col is not None and duration_col is not None and not se_df.empty: + for _, row in se_df.iterrows(): + end_ts = row[current_time_col] + dur = row[duration_col] + if pd.isna(end_ts): + continue + if pd.notna(dur) and float(dur) > 0: + start_ts = end_ts - pd.to_timedelta(float(dur), unit="ms") + else: + start_ts = end_ts + if end_ts >= start_ts: + intervals.append((start_ts, end_ts)) + + se_ms = 0.0 + if intervals: + ordered = sorted(intervals, key=lambda x: (x[0], x[1])) + se_total_ms = 0.0 + cur_start, cur_end = ordered[0] + for start, end in ordered[1:]: + if start <= cur_end: + if end > cur_end: + cur_end = end + else: + se_total_ms += (cur_end - cur_start).total_seconds() * 1000.0 + cur_start, cur_end = start, end + se_total_ms += (cur_end - cur_start).total_seconds() * 1000.0 + se_ms = round(se_total_ms, 3) + elif duration_col is not None and not se_df.empty: + se_total_ms = pd.to_numeric(se_df[duration_col], errors="coerce").fillna(0.0).clip(lower=0.0).sum() + if se_total_ms > 0: + se_ms = round(min(float(se_total_ms), total_ms), 3) + + fe_ms = round(max(0.0, total_ms - se_ms), 3) + + summary = { + "Dataset": dataset, + "Workspace": workspace, + "Total Elapsed (ms)": total_ms, + "Storage Engine (ms)": se_ms, + "Formula Engine (ms)": fe_ms, + "Trace Event Count": int(len(trace_df)), + "Note": "Timings are derived from QueryEnd + SE end events (VertiPaq/DirectQuery).", + } + + return pd.DataFrame([summary]) + + @log def _dax_perf_test( dataset: str, diff --git a/tests/test_capture_dax_query_timings.py b/tests/test_capture_dax_query_timings.py new file mode 100644 index 000000000..ec24778ac --- /dev/null +++ b/tests/test_capture_dax_query_timings.py @@ -0,0 +1,161 @@ +import pandas as pd + +import sempy_labs._dax as dax_module + + +class _FakeTrace: + def __init__(self, df): + self._df = df + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def start(self): + return None + + def stop(self): + return self._df + + +class _FakeTraceConnection: + def __init__(self, df): + self._df = df + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def create_trace(self, event_schema): + return _FakeTrace(self._df) + + +def test_capture_dax_query_timings_returns_summary_df(monkeypatch): + run_df = pd.DataFrame( + [ + { + "EventClass": "VertiPaqSEQueryEnd", + "CurrentTime": "2026-05-25T12:00:00.200", + "Duration": 120, + "SessionID": "s1", + "ActivityID": "a1", + "RequestID": "r1", + "TextData": "SE event", + }, + { + "EventClass": "QueryEnd", + "CurrentTime": "2026-05-25T12:00:00.300", + "Duration": 300, + "SessionID": "s1", + "ActivityID": "a1", + "RequestID": "r1", + "TextData": "-- RUNID:abc EVALUATE ROW(\"x\", 1)", + }, + ] + ) + + monkeypatch.setattr(dax_module, "uuid4", lambda: type("U", (), {"hex": "abc"})()) + monkeypatch.setattr(dax_module.time, "sleep", lambda _: None) + monkeypatch.setattr(dax_module.fabric, "create_trace_connection", lambda dataset, workspace=None: _FakeTraceConnection(run_df)) + captured = {} + + def _fake_evaluate_dax(dataset, workspace, dax_string): + captured["dax_string"] = dax_string + return pd.DataFrame() + + monkeypatch.setattr(dax_module.fabric, "evaluate_dax", _fake_evaluate_dax) + + calls = {"clear": 0} + + def _fake_clear_cache(dataset, workspace=None): + calls["clear"] += 1 + + monkeypatch.setattr(dax_module, "clear_cache", _fake_clear_cache, raising=False) + monkeypatch.setattr("sempy_labs._clear_cache.clear_cache", _fake_clear_cache) + + result = dax_module.capture_dax_query_timings(dataset="Model", dax_query='EVALUATE ROW("x", 1)', workspace="WS") + + assert isinstance(result, pd.DataFrame) + assert len(result) == 1 + assert result.loc[0, "Dataset"] == "Model" + assert result.loc[0, "Workspace"] == "WS" + assert result.loc[0, "Total Elapsed (ms)"] == 300.0 + assert result.loc[0, "Storage Engine (ms)"] == 120.0 + assert result.loc[0, "Formula Engine (ms)"] == 180.0 + assert result.loc[0, "Trace Event Count"] == 2 + assert calls["clear"] == 1 + assert captured["dax_string"] == '-- RUNID:abc\nEVALUATE ROW("x", 1)' + assert captured["dax_string"].startswith("-- RUNID:abc\n") + assert 'EVALUATE ROW("x", 1)' in captured["dax_string"] + + +def test_capture_dax_query_timings_empty_trace(monkeypatch): + empty_df = pd.DataFrame() + + monkeypatch.setattr(dax_module.time, "sleep", lambda _: None) + monkeypatch.setattr(dax_module.fabric, "create_trace_connection", lambda dataset, workspace=None: _FakeTraceConnection(empty_df)) + monkeypatch.setattr(dax_module.fabric, "evaluate_dax", lambda dataset, workspace, dax_string: pd.DataFrame()) + + perf_values = iter([1.0, 1.1, 1.5]) + monkeypatch.setattr(dax_module.time, "perf_counter", lambda: next(perf_values)) + + result = dax_module.capture_dax_query_timings( + dataset="Model", + dax_query='EVALUATE ROW("x", 1)', + clear_cache=False, + ) + + assert isinstance(result, pd.DataFrame) + assert len(result) == 1 + assert result.loc[0, "Total Elapsed (ms)"] == 100.0 + assert result.loc[0, "Storage Engine (ms)"] == 0.0 + assert result.loc[0, "Trace Event Count"] == 0 + + +def test_capture_dax_query_timings_se_duration_fallback(monkeypatch): + fallback_df = pd.DataFrame( + [ + { + "EventClass": "VertiPaqSEQueryEnd", + "CurrentTime": "not-a-timestamp", + "Duration": 120, + "SessionID": "s1", + "ActivityID": "a1", + "RequestID": "r1", + "TextData": "SE event", + }, + { + "EventClass": "QueryEnd", + "CurrentTime": "not-a-timestamp", + "Duration": 300, + "SessionID": "s1", + "ActivityID": "a1", + "RequestID": "r1", + "TextData": "-- RUNID:abc EVALUATE ROW(\"x\", 1)", + }, + ] + ) + + monkeypatch.setattr(dax_module.time, "sleep", lambda _: None) + monkeypatch.setattr(dax_module.fabric, "create_trace_connection", lambda dataset, workspace=None: _FakeTraceConnection(fallback_df)) + monkeypatch.setattr(dax_module.fabric, "evaluate_dax", lambda dataset, workspace, dax_string: pd.DataFrame()) + monkeypatch.setattr(dax_module, "uuid4", lambda: type("U", (), {"hex": "abc"})()) + + perf_values = iter([1.0, 1.1, 1.4]) + monkeypatch.setattr(dax_module.time, "perf_counter", lambda: next(perf_values)) + + result = dax_module.capture_dax_query_timings( + dataset="Model", + dax_query='EVALUATE ROW("x", 1)', + clear_cache=False, + ) + + assert isinstance(result, pd.DataFrame) + assert len(result) == 1 + assert result.loc[0, "Total Elapsed (ms)"] == 300.0 + assert result.loc[0, "Storage Engine (ms)"] == 120.0 + assert result.loc[0, "Formula Engine (ms)"] == 180.0