Skip to content

fix(intake): performance improvements for trace list #1002

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BrianNewsom merged 2 commits into from
Jul 30, 2026
Merged

fix(intake): performance improvements for trace list #1002

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4b541a5
fix(intake): performance improvements for trace list
BrianNewsom Jul 30, 2026
27ea1b3
fix(intake): reconcile trace page hydration
BrianNewsom Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
303 changes: 193 additions & 110 deletions services/intake/src/nmp/intake/repository/clickhouse/trace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
from __future__ import annotations

import json
import logging
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

Expand All @@ -25,6 +28,8 @@
text_select_for_mode,
)

logger = logging.getLogger(__name__)

TRACE_SORT_COLUMNS = {
"started_at": "started_at",
}
Expand Down Expand Up @@ -77,6 +82,23 @@
_ZERO_DATETIME = datetime.fromtimestamp(0, tz=timezone.utc)


@dataclass(frozen=True)
class _TracePageRef:
source_format: str
trace_id: str

@classmethod
def from_row(cls, row: Mapping[str, Any]) -> _TracePageRef:
return cls(
source_format=str(row["source_format"]),
trace_id=str(row["id"]),
)

@property
def trace_key(self) -> tuple[str, str]:
return self.source_format, self.trace_id


class ClickHouseTraceRepository(TraceRepository):
def __init__(self, executor: ClickHouseExecutor) -> None:
self._executor = executor
Expand Down Expand Up @@ -113,29 +135,48 @@ async def list_traces(
)

offset = (page - 1) * page_size
trace_index_sql, row_parameters = _trace_index_sql(
trace_index_table,
filters,
mode=mode,
)
rows_sql, rows_parameters = _trace_rows_sql(
trace_index_sql=trace_index_sql,
spans_table=spans_table,
mode=mode,
sort=sort,
)
rows = await self._executor.fetch_all(
page_sql = _trace_page_sql(trace_index_sql=count_trace_index_sql, sort=sort)
page_rows = await self._executor.fetch_all(
ClickHouseQuery(
name="traces.list.rows",
statement=rows_sql,
name="traces.list.page",
statement=page_sql,
parameters={
**row_parameters,
**rows_parameters,
**parameters,
"limit": page_size,
"offset": offset,
},
)
)
rows = page_rows
if mode != "summary" and page_rows:
page_refs = [_TracePageRef.from_row(row) for row in page_rows]
hydration_sql, hydration_parameters = _trace_hydration_sql(
trace_index_table=trace_index_table,
spans_table=spans_table,
mode=mode,
)
rows = await self._executor.fetch_all(
ClickHouseQuery(
name="traces.list.hydrate",
statement=hydration_sql,
parameters={
"workspace": filters.workspace,
**hydration_parameters,
**_trace_page_parameters(page_refs),
},
)
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
rows, dropped_refs = _reconcile_hydrated_page(page_refs, rows)
if dropped_refs:
logger.warning(
"Trace page hydration omitted refs returned by the page query",
extra={
"workspace": filters.workspace,
"page": page,
"dropped_trace_refs": [ref.trace_key for ref in dropped_refs],
},
)
total_results = max(offset + len(rows), total_results - len(dropped_refs))
traces = [_row_to_trace(row) for row in rows]
return PaginatedResult(
data=traces,
Expand Down Expand Up @@ -212,47 +253,59 @@ async def latest_trace_started_at_by_group(
return {str(row["group_id"]): row["started_at"] for row in rows}


def _trace_rows_sql(
*, trace_index_sql: str, spans_table: str, mode: TraceMode, sort: str
) -> tuple[str, dict[str, Any]]:
if mode == "summary":
query = f"""
WITH page_traces AS (
SELECT *
FROM ({trace_index_sql}) AS traces
ORDER BY {_order_by(sort)}
LIMIT %(limit)s OFFSET %(offset)s
)
SELECT
{_trace_select_columns(include_aggregates=False)}
FROM page_traces AS traces
ORDER BY {_order_by(sort, table_alias="traces")}
"""
return query, {}
if mode in {"preview", "detailed"}:
aggregates_sql, parameters = _trace_aggregates_sql(spans_table)
query = f"""
WITH
page_traces AS (
SELECT *
FROM ({trace_index_sql}) AS traces
ORDER BY {_order_by(sort)}
LIMIT %(limit)s OFFSET %(offset)s
),
rollups AS (
{aggregates_sql}
)
SELECT
{_trace_select_columns(include_aggregates=True)}
FROM page_traces AS traces
LEFT JOIN rollups
ON traces.workspace = rollups.workspace
AND traces.source_format = rollups.source_format
AND traces.id = rollups.trace_id
ORDER BY {_order_by(sort, table_alias="traces")}
"""
return query, parameters
raise ValueError(f"Unsupported trace mode: {mode}")
def _trace_page_sql(*, trace_index_sql: str, sort: str) -> str:
"""Build the first-phase query that selects lightweight trace roots for one page."""

return f"""
SELECT
{_trace_select_columns(include_aggregates=False)}
Comment thread
BrianNewsom marked this conversation as resolved.
FROM ({trace_index_sql}) AS traces
ORDER BY {_order_by(sort, table_alias="traces")}
LIMIT %(limit)s OFFSET %(offset)s
"""


def _trace_hydration_sql(*, trace_index_table: str, spans_table: str, mode: TraceMode) -> tuple[str, dict[str, Any]]:
"""Build the second-phase query that hydrates page roots and span aggregates."""

trace_roots_sql, trace_parameters = _page_trace_roots_sql(trace_index_table=trace_index_table, mode=mode)
aggregates_sql, aggregate_parameters = _trace_aggregates_sql(spans_table)
query = f"""
Comment thread
BrianNewsom marked this conversation as resolved.
WITH
traces AS (
{trace_roots_sql}
),
rollups AS (
{aggregates_sql}
)
SELECT
{_trace_select_columns(include_aggregates=True)}
FROM traces
LEFT JOIN rollups
ON traces.workspace = rollups.workspace
AND traces.source_format = rollups.source_format
AND traces.id = rollups.trace_id
"""
return query, {**trace_parameters, **aggregate_parameters}


def _page_trace_roots_sql(*, trace_index_table: str, mode: TraceMode) -> tuple[str, dict[str, Any]]:
trace_columns, parameters = _trace_index_select_columns(mode=mode)
query = f"""
SELECT
{trace_columns}
FROM {trace_index_table} AS trace_roots FINAL
WHERE trace_roots.workspace = %(workspace)s
AND trace_roots.is_deleted = 0
AND trace_roots.trace_id IN %(page_trace_ids)s
AND (
trace_roots.source_format,
trace_roots.trace_id
) IN %(page_trace_keys)s
ORDER BY trace_roots.root_started_at ASC, trace_roots.root_span_id ASC
LIMIT 1 BY trace_roots.workspace, trace_roots.source_format, trace_roots.trace_id
"""
return query, parameters


def _trace_select_columns(*, include_aggregates: bool) -> str:
Expand Down Expand Up @@ -294,29 +347,11 @@ def _trace_index_sql(
mode: TraceMode,
) -> tuple[str, dict[str, Any]]:
where_sql, parameters = _trace_index_where(filters, qualifier="trace_roots")
parameters.update(text_query_parameters(mode))
payload_columns = ",\n ".join(
(
text_select_for_mode("trace_roots.root_input", alias="input", mode=mode),
text_select_for_mode("trace_roots.root_output", alias="output", mode=mode),
)
)
select_columns, select_parameters = _trace_index_select_columns(mode=mode)
parameters.update(select_parameters)
query = f"""
SELECT
trace_roots.trace_id AS id,
trace_roots.workspace AS workspace,
trace_roots.session_id AS session_id,
trace_roots.source_format AS source_format,
nullIf(trace_roots.root_span_id, '') AS root_span_id,
nullIf(trace_roots.root_name, '') AS name,
{payload_columns},
nullIf(trace_roots.project, '') AS project,
nullIf(trace_roots.evaluation_id, '') AS evaluation_id,
nullIf(trace_roots.test_case_id, '') AS test_case_id,
trace_roots.root_started_at AS started_at,
trace_roots.root_ended_at AS ended_at,
trace_roots.root_status AS status,
trace_roots.event_ts AS ingested_at
{select_columns}
FROM {table} AS trace_roots FINAL
WHERE {where_sql}
ORDER BY trace_roots.root_started_at ASC, trace_roots.root_span_id ASC
Expand All @@ -325,48 +360,96 @@ def _trace_index_sql(
return query, parameters


def _trace_aggregates_sql(table: str) -> tuple[str, dict[str, Any]]:
source_alias = "trace_spans"
metric_columns, parameters = metric_aggregate_columns(source_alias)
def _trace_index_select_columns(*, mode: TraceMode) -> tuple[str, dict[str, Any]]:
payload_columns = (
text_select_for_mode("trace_roots.root_input", alias="input", mode=mode),
text_select_for_mode("trace_roots.root_output", alias="output", mode=mode),
)
columns = (
"trace_roots.trace_id AS id",
"trace_roots.workspace AS workspace",
"trace_roots.session_id AS session_id",
"trace_roots.source_format AS source_format",
"nullIf(trace_roots.root_span_id, '') AS root_span_id",
"nullIf(trace_roots.root_name, '') AS name",
*payload_columns,
"nullIf(trace_roots.project, '') AS project",
"nullIf(trace_roots.evaluation_id, '') AS evaluation_id",
"nullIf(trace_roots.test_case_id, '') AS test_case_id",
"trace_roots.root_started_at AS started_at",
"trace_roots.root_ended_at AS ended_at",
"trace_roots.root_status AS status",
"trace_roots.event_ts AS ingested_at",
)
return ",\n ".join(columns), text_query_parameters(mode)

model_spec = spec_for_field(SpanAttributeField.MODEL)
provider_spec = spec_for_field(SpanAttributeField.PROVIDER)
parameters["model_key"] = model_spec.bag_key
parameters["provider_key"] = provider_spec.bag_key

def _trace_aggregates_sql(table: str) -> tuple[str, dict[str, Any]]:
source_alias = "trace_spans"
select_columns, parameters = _trace_aggregate_select_columns(source_alias)
current_spans = current_spans_sql(
table,
extra_where_sql=(
"span_versions.trace_id IN %(page_trace_ids)s\n"
" AND (span_versions.source_format, span_versions.trace_id) "
"IN %(page_trace_keys)s"
),
)
query = f"""
SELECT
{source_alias}.workspace AS workspace,
{source_alias}.source_format AS source_format,
{source_alias}.trace_id AS trace_id,
{metric_columns},
arraySort(groupUniqArrayIf(
{source_alias}.attributes_string[%(model_key)s],
has(mapKeys({source_alias}.attributes_string), %(model_key)s)
AND {source_alias}.attributes_string[%(model_key)s] != ''
)) AS models,
arraySort(groupUniqArrayIf(
{source_alias}.attributes_string[%(provider_key)s],
has(mapKeys({source_alias}.attributes_string), %(provider_key)s)
AND {source_alias}.attributes_string[%(provider_key)s] != ''
)) AS providers,
count() AS span_count,
countIf({source_alias}.status = 'error') AS error_count
FROM {
current_spans_sql(
table,
extra_where_sql=(
"(span_versions.workspace, span_versions.source_format, span_versions.trace_id) IN "
"(SELECT workspace, source_format, id FROM page_traces)"
),
)
} AS {source_alias}
{select_columns}
FROM {current_spans} AS {source_alias}
WHERE {source_alias}.is_deleted = 0
GROUP BY {source_alias}.workspace, {source_alias}.source_format, {source_alias}.trace_id
"""
return query, parameters


def _trace_aggregate_select_columns(source_alias: str) -> tuple[str, dict[str, Any]]:
metric_columns, parameters = metric_aggregate_columns(source_alias)
parameters.update(
model_key=spec_for_field(SpanAttributeField.MODEL).bag_key,
provider_key=spec_for_field(SpanAttributeField.PROVIDER).bag_key,
)
columns = (
f"{source_alias}.workspace AS workspace",
f"{source_alias}.source_format AS source_format",
f"{source_alias}.trace_id AS trace_id",
metric_columns,
_unique_string_attribute_aggregate(source_alias, parameter="model_key", alias="models"),
_unique_string_attribute_aggregate(source_alias, parameter="provider_key", alias="providers"),
"count() AS span_count",
f"countIf({source_alias}.status = 'error') AS error_count",
)
return ",\n ".join(columns), parameters


def _unique_string_attribute_aggregate(source_alias: str, *, parameter: str, alias: str) -> str:
attribute = f"{source_alias}.attributes_string[%({parameter})s]"
return f"""arraySort(groupUniqArrayIf(
{attribute},
has(mapKeys({source_alias}.attributes_string), %({parameter})s)
AND {attribute} != ''
)) AS {alias}"""


def _trace_page_parameters(refs: Sequence[_TracePageRef]) -> dict[str, object]:
return {
"page_trace_ids": [ref.trace_id for ref in refs],
"page_trace_keys": [ref.trace_key for ref in refs],
}


def _reconcile_hydrated_page(
page_refs: Sequence[_TracePageRef],
hydrated_rows: Sequence[dict[str, Any]],
) -> tuple[list[dict[str, Any]], list[_TracePageRef]]:
rows_by_ref = {_TracePageRef.from_row(row): row for row in hydrated_rows}
rows = [rows_by_ref[ref] for ref in page_refs if ref in rows_by_ref]
dropped_refs = [ref for ref in page_refs if ref not in rows_by_ref]
return rows, dropped_refs


# Maps API/filter field names to their physical trace_index columns.
_TRACE_INDEX_FILTER_COLUMNS = {
"evaluation_id": "evaluation_id",
Expand Down
Loading
Loading