Add a RFC to optimize common trace DB queries by mprahl · Pull Request #24 · mlflow/rfcs

mprahl · 2026-07-08T20:33:14Z

Supersedes #23

This was coauthored by @aws-khatria.

Co-authored-by: Ayush Khatri <khatria@amazon.com> Signed-off-by: mprahl <mprahl@users.noreply.github.com>

etirelli · 2026-07-08T21:05:01Z

+    ON assessments (experiment_id, trace_timestamp_ms);
+
+CREATE INDEX idx_assessments_exp_trace_ts_name
+    ON assessments (experiment_id, trace_timestamp_ms, name);


aren't the two indexes above redundant? I believe in postgres and SQLite, the second index can be used for any query the first would.

I think that it can be reused but it's not as efficient. I added a comment explaining it.

etirelli · 2026-07-08T21:10:09Z

+
+These numbers are benchmark-based targets, not guarantees. Exact percentile queries may still
+retain significant sort cost even after the denormalization work, and are expected to remain the
+main residual cost on the slowest paths.


I think it would be good to also benchmark writes

etirelli

LGTM, just a couple of minor comments.

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl · 2026-07-09T21:49:42Z

@etirelli could you please take another look? I added a commit with a proposal for rollup tables.

jeremyeder · 2026-07-09T23:15:01Z

+- adding targeted indexes for assessment and span-cost workloads
+- cleaning up duplicated EAV rows once the denormalized columns become authoritative
+
+On a load test with 10M traces, 30M spans, and 30M assessments in one experiment, the current Usage


hey, is this load generator available somewhere?

There are some load generators available in the MLflow repo:

.github/workflows/tracing-benchmark.yml

dev/benchmarks/tracing/README.md

dev/benchmarks/tracing/test_trace_perf.py

dev/benchmarks/tracing/_data.py

That said, the exact 10M-trace Postgres UI replay harness used for these RFC numbers was a local/ad hoc load harness derived from the tracing benchmark utilities to speed up trace ingestion.

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

etirelli · 2026-07-10T13:50:51Z

+similar non-analytic annotations can remain mutable. Adding assessments to older traces should also
+remain allowed, since human review and evaluator backfills often happen after the trace is closed.
+
+Assessment writes or trace deletes against already-rolled-up traces should delete stale rollup rows


should we mark the row as stale instead of deleting? considering queries to the metrics issued after an update to a non-fresh trace, if the row was deleted, the data point for that day would be missing completely, while if we simply mark it as stale, the previous result would be returned/computed (eventual consistency).

@etirelli I'm torn about this. I think I lean towards the eventual consistency you suggest but it does mean the metrics would be inaccurate. Perhaps a compromise is to show a warning in the UI that the metrics are out of date and being refreshed?

What do you think @B-Step62?

etirelli · 2026-07-10T13:54:45Z

+rollup-supported shape, especially arbitrary filtered requests and non-daily exact percentile
+queries, still fall back to the raw path and remain the main residual cost.
+
+## Drawbacks


NP: "trade-offs" instead of "drawbacks"? :)

etirelli · 2026-07-10T13:57:54Z

@mprahl LGTM

The query numbers look great!

Did you benchmark writing impact on a loaded DB (i.e., writing performance after pre-loading the DB with the test data)? I don't believe the impact would be significant, but still a good datapoint to have.

mprahl · 2026-07-10T14:24:48Z

@mprahl LGTM

The query numbers look great!

Did you benchmark writing impact on a loaded DB (i.e., writing performance after pre-loading the DB with the test data)? I don't believe the impact would be significant, but still a good datapoint to have.

Thanks for the review! I haven't tested it but I'll try to give that a shot today.

kramaranya · 2026-07-13T13:40:47Z

+| Full UI replay, 20 serial requests | 1,891.9s        | 119.4s                    | 27.5s                  | 68.8x                    |
+| Cost over time by model            | 110.3s          | 21.3s                     | 59ms                   | 1,870x                   |
+| Cost breakdown by model            | 24.2s           | 14.4s                     | 52ms                   | 464x                     |
+| Token/latency daily time series    | 9.4-459.0s      | 5.4-9.8s                  | 56-63ms                | 159-7,908x               |
+| Time-range trace count             | 315ms           | 558ms                     | 53ms                   | 5.9x                     |


kramaranya · 2026-07-13T13:49:09Z

+To keep rollups stable, MLflow should reject late mutations that change rollup facts once a trace is
+older than the freshness window. Specifically, the MLflow API should not allow adding spans, adding
+trace metrics, or updating trace-info fields such as timestamp, duration, status, token usage, span
+cost, or other analytic facts for traces older than 24 hours. Tags, comments, display metadata, and


Is this 24-hour immutability rule scoped to rollup-enabled deployments only?

Good question, I was thinking it would apply to all deployments by default but would welcome opinions on this.

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

HumairAK · 2026-07-13T20:01:51Z

+
+1. compute eligible `(workspace, experiment_id, rollup_day)` partitions from raw traces
+2. exclude days newer than a freshness threshold, defaulting to 24 hours
+3. rebuild deleted/stale rollup partitions before processing newly eligible trace days


One obvious race is:

rollup job starts rebuilding day D

it reads raw rows for D

meanwhile an assessment is added to an old trace in D

invalidation logic deletes the old rollup rows for D

job finishes and writes rollups based on the old snapshot, missing that new assessment

That can leave you with a day that looks "complete" again but is already stale.

Can we add a generation/version number to the partigion generation, or have some other means of addressing such race conditions?

Other potential risks:

two rebuilders at once: two job workers rebuild the same partition concurrently; with deterministic upserts this may be mostly duplicate work, but last-writer-wins can still publish an older snapshot

partial visibility: if rows are inserted piecemeal, a reader could see half a partition unless publication is atomic

marker-row false confidence: if trace_count is used only as "exists => done", a stale or partially rebuilt partition can look valid

delete vs rebuild race: invalidation deletes rows while a builder is repopulating them, producing timing-dependent results

For two rebuilders - I think we can just limit it to one job at a time instead of supporting multipler workers, unless you feel like adressing it.

I updated the RFC to make it clear that this is serial. The implication is this is just a Huey job under the hood like trace archival.

HumairAK · 2026-07-13T21:26:14Z

+
+`experiment_id` should remain denormalized but the foreign key constraint removed. MLflow already
+knows the experiment from the owning trace when it writes the assessment, so storing that value does
+not require an extra integrity lookup on every insert. Skipping the foreign key keeps this


Dropping the FK feels like it could lead to data corruption if application logic fails. Is the gain in lookup here enough to justify the risk?

If so then as a precaution we should just always ensure we derive assessments.experiment_id inside the store from the owning trace_info and never from the client payload.

The foreign key removal isn't for improving reads but is for improving writes in order for the DB to skip integrity verification on every trace insertion. I can drop it from the proposal since this RFC is about improving read performance, however, I do think there is benefit here.

Sorry yeah I meant for writes. The benefit I think has to justify the added risk to integrity, I'll leave the judgement call to you, I lean towards keeping the constraint.

HumairAK · 2026-07-13T22:09:58Z

+
+The span-cost rollup table stores daily `input_cost`, `output_cost`, and `total_cost` aggregates by
+model and provider grains. The assessment rollup table stores daily assessment count and numeric
+assessment-value aggregates by assessment name and value grains. The rollup reader should use these


Probably too technical for this RFC but you mention Rollup Reader, I think some light abstraction might be warranted given that other SQL db's seem to have similar concepts, and we don't just bake into the logic something like:

if postgres: use MATERIALIZED CTE
if postgres: use rollup tables
if postgres: use these invalidation rules
if postgres: use these indexes

--

So maybe something like a planner with a postgresql implementation:

Where a generic interface / planner answers:

“Can this request use rollups?”

“Execute raw path or rollup path”

“Invalidate this partition”

And Postgres implementation would:

exact SQL

exact tables

exact indexes

exact maintenance job behavior

I added some high level text about the query planner should abstract these decisions and not at the DB query level.

HumairAK · 2026-07-13T22:25:45Z

+| -------------------------------- | ------------------------------------------------------------------------------------------ | ------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- |
+| `sql_trace_metric_daily_rollups` | `(workspace, experiment_id, rollup_day, metric_name, trace_status?, trace_name?)`          | `trace_status`, `trace_name`                                        | `sample_count`, `sum_value`, `min_value`, `max_value`, `p50_value`, `p90_value`, `p99_value` |
+| `sql_span_cost_daily_rollups`    | `(workspace, experiment_id, rollup_day, metric_name, model_name?, model_provider?)`        | `model_name`, `model_provider`                                      | `sample_count`, `sum_value`, `min_value`, `max_value`                                        |
+| `sql_assessment_daily_rollups`   | `(workspace, experiment_id, rollup_day, metric_name, assessment_name?, assessment_value?)` | `assessment_name`, `assessment_value_json`, `assessment_value_text` | `sample_count`, `sum_value`, `min_value`, `max_value`                                        |


Based on quick discussions with GPT, it's my understanding that Rollups work best when grouping dimensions are stable and bounded. This doesn't seem like it would be the case with user-defined assessment_value, which can take numeric scores like 0.82.

Would it be better if instead we didn't roll up by default on assessment_value, or roll up only for eligible low-cardinality types like boolean or enum-like strings?

Very good point! I removed this. Sadly, it made the benchmark about 8 seconds slower which prompted me to add a configuration to cap the max assessment summary when clicking on the Traces tab.

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

HumairAK · 2026-07-14T20:25:09Z

+5. rebuild a bounded number of missing partitions per pass
+6. upsert rollup rows by deterministic rollup key
+
+Invalidated partitions should be tracked durably until a rebuild completes. This ensures deletes or


Thanks, for this part:

"Invalidated partitions should be tracked durably until a rebuild completes. This ensures deletes or moves can repair an old partition even when no source rows remain there."

I think we should clarify what we mean by "durably". Do you mean persist in DB somewhere?

B-Step62 · 2026-07-14T19:46:45Z

+They are distinct from span cost fields, which support model/provider analytics. Both should be
+backfilled and kept consistent by their respective write paths before legacy cost rows are removed.
+
+`trace_name` should also become the authoritative storage and query field on `trace_info` for


What is the query pattern benefit from denormalizing trace name? Or is it mainly for reduce one tag row write at ingestion?

The primary benefit is that filtering, sorting, and grouping by trace name can use trace_info directly instead of joining trace_tags. These queries are not part of the default UI path, so this is not a major contributor to the benchmark improvement. However, trace name is a first-class trace attribute, so storing it as a dedicated column also gives us a cleaner schema.

I can keep it or drop it. Do you have a preference?

B-Step62 · 2026-07-14T19:48:38Z

+
+### 2. Denormalize hot span cost analytics onto `spans`
+
+Add the following columns to `spans`:


Span costs are used for cost breakdown per model/provider, but the model and provider is still stored in attributes IIRC, do we need to denormalize them as well?

Good point! I'll update the RFC to denormalize them and drop the dimension_attributes column as it only held those two values.

B-Step62 · 2026-07-14T19:53:49Z

+GROUP BY name;
+```
+
+The optimized query aggregates directly on a numeric column. Without `aggregate_value`, the hot path


Q: I guess the aggregate_value also simplify the query for filtering traces based on assessment value? Right now we need a suboptimal query for numeric assessment filtering because the assessment raw value stored in the table is string, it's nice if we can clean it up too.

Good point. I added an is_numeric_value column to allow using aggregate_value on numeric comparison queries.

B-Step62 · 2026-07-14T20:08:07Z

+
+Add three daily rollup tables:
+
+| Table                            | Grain                                                                               | Dimension columns              | Measure columns                                                                              |


Have we considered using higher granularity like hour for rollup table rows? MLflow's built-in dashboard support hourly as the smallest time unit.

One downside tho is p50/p90/p99 handling for latency metric becomes tricker - because that's not aggregatable to bigger unit unlike sum, average, min/max.

Have we considered using higher granularity like hour for rollup table rows?

I think the UI only uses hourly buckets when the time range is 7 days or less, so since that date range is on the smaller end, I'd tend to favor doing raw queries rather than store hourly rollups. The performance for that isn't too bad after the column denormalization.

What do you think?

B-Step62 · 2026-07-14T20:14:45Z

+dimensions or aggregations, unsupported percentile values, or bucket widths other than one day
+should fall back to the raw query path.
+
+The Traces table currently requests exact assessment-value counts when infinite pagination is


Infinite pagination is generally turned off for all users. It is gated by a feature flag in UI code, not server configuration users enable, so I don't think we need this configuration.

The only place it still uses an infinite scroll is RunViewEvaluationsTab.tsx, but I doubt users have eval runs with million of rows.

Sorry, this was misleading. The overall assessment value summary is for the whole date range on the "Traces" tab, not just the the current page. I'll add an open question about whether or not to add a configuration to cap this.

B-Step62 · 2026-07-14T20:19:23Z

+
+1. compute eligible `(workspace, experiment_id, rollup_day)` partitions from raw trace, span, and
+   assessment rows
+2. exclude days newer than a configurable freshness threshold, controlled by


I'm not a big fan of making many things configurable upfront. Do we have potential use cases for server admins to change this freshness value, and is there any guide we will provide for them to pick the right vlaue? If not, I would start with a fixed value.

What I think is more useful for admins is a configuration to decide the time of day to trigger this rollup. In large scale this job can be a heavy workload, then admins might want to pick a low-traffic time to trigger this job.

I dropped the freshness threshold configuration and also reworded the RFC to allow old spans and traces to be ingested, it just causes the rollup day to be marked stale to be rebuilt.

I also made the rollup configuration be configurable via a cron schedule.

B-Step62

Left a few comments on rollup correctness and cross-backend behavior. All comments generated by Claude (reviewed by me before posting).

B-Step62 · 2026-07-14T20:27:27Z

+- `total_cost`
+
+Span cost aggregation should read these columns directly. Any remaining cost readers, including
+`sum_gateway_trace_cost()`, should also be retargeted to `spans.total_cost` so cost metrics do not


Section 1 says the trace-level cost fields on trace_info are authoritative for gateway readers, but this section retargets sum_gateway_trace_cost() to spans.total_cost. Which one is authoritative? Note the current query is a workspace-wide time-window sum with no experiment filter, so neither proposed span-cost index (both lead with trace_id / experiment_id) covers it, while reading trace_info.total_cost would drop one join from the budget enforcement hot path.

Generated by Claude

Clarified in the RFC. spans costs are used for model/provider and total_cost are used by the Gateway and trace level queries.

B-Step62 · 2026-07-14T20:27:27Z

+
+When SQL rollups are enabled, to keep rollups stable, MLflow should reject late mutations that
+change rollup facts once a trace is older than the freshness window. Specifically, the MLflow API
+should not allow adding spans, adding trace metrics, or updating trace-info fields such as


This list covers late mutations but not late creation: a trace inserted with a backdated timestamp_ms (delayed OTel batch export, bulk import of historical traces) lands in a day that already has a coverage marker and silently never appears in rollups. Should backdated trace creation invalidate the partition instead? Rejecting it like the other late mutations would break import tooling.

Generated by Claude

Dropped the proposal to reject late coming traces and just cause the rollup values to be queued for rebuilds.

B-Step62 · 2026-07-14T20:27:27Z

+  columns to match the existing analytics numeric path.
+- In the Alembic migration code, `batch_op` should be used only for SQLite compatibility paths.
+  Large Postgres or MySQL tables should use direct `ALTER TABLE` / `CREATE INDEX` statements to
+  avoid expensive table recreation.


Since the Alembic migration line is shared across backends, can we spell out the story for MySQL / MSSQL / SQLite? The proposed indexes use INCLUDE + partial WHERE (MySQL supports neither, SQLite lacks INCLUDE), while the denormalized columns, read retargeting, and EAV cleanup apply to every backend. It would be good to state what each dialect gets and whether the rollup tables/jobs are Postgres-only.

Generated by Claude

The RFC wording was ambiguous, the intention is for all DB backends to be supported, however, there will be specific Postgres optimizations (e.g. INCLUDE on one of the indexes) and the percentile queries are mostly Postgres only. Updated the RFC.

B-Step62 · 2026-07-14T20:27:27Z

+should be centralized behind backend-specific planning code rather than scattered across individual
+query paths. Requests with arbitrary filters, exact assessment-value dimensions, unsupported
+dimensions or aggregations, unsupported percentile values, or bucket widths other than one day
+should fall back to the raw query path.


The UI sends arbitrary start_time_ms / end_time_ms (typically now minus N days), so the leading bucket of nearly every real request is a partial UTC day, and serving it from a daily rollup would over-count. The freshness window handles the trailing edge, but can we state explicitly that partial edge days must be computed from raw rows and stitched with rollup days? The fixed-range benchmark replays would not have exercised this path.

Generated by Claude

Added clarification that the first and last non-full days are calculated raw and the rest use the rollups.

B-Step62 · 2026-07-14T20:27:27Z

+1. backfill `trace_info.trace_name` from trace tags
+2. retarget trace-name readers, including trace-name analytics filters and rollup builders, to
+   `trace_info.trace_name`
+3. synthesize the API-facing trace-name tag from `trace_info.trace_name` for compatibility


Synthesis covers reads, but what happens when a client writes these keys directly after migration, e.g. SetTraceTag / DeleteTraceTag on mlflow.traceName, metadata updates on mlflow.trace.session, or logging a trace metric with one of the reserved token keys? Without redirect-to-column (or rejection) semantics the columns and the synthesized API views can drift. Also worth stating explicitly that search_traces filter strings like tags.mlflow.traceName = '...' are retargeted too.

Generated by Claude

Clarified that the API remains compatible.

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

Add a RFC to optimize common trace DB queries

b18b2ec

Co-authored-by: Ayush Khatri <khatria@amazon.com> Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl mentioned this pull request Jul 8, 2026

RFC: PostgreSQL trace query optimization #23

Closed

etirelli reviewed Jul 8, 2026

View reviewed changes

etirelli approved these changes Jul 8, 2026

View reviewed changes

Add a proposal for rollup tables

a0946bf

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl requested a review from etirelli July 9, 2026 21:49

jeremyeder reviewed Jul 9, 2026

View reviewed changes

Add baseline benchmark

bae7ca2

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

etirelli reviewed Jul 10, 2026

View reviewed changes

etirelli approved these changes Jul 10, 2026

View reviewed changes

kramaranya reviewed Jul 13, 2026

View reviewed changes

mprahl added 2 commits July 13, 2026 13:07

Address feedback and inconsistencies

e5f5594

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

Add the ability to preload denormalized columns to reduce downtime

4fe9e8c

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

B-Step62 assigned PattaraS and unassigned PattaraS Jul 13, 2026

B-Step62 requested review from B-Step62 and PattaraS July 13, 2026 20:11

HumairAK reviewed Jul 13, 2026

View reviewed changes

Address feedback

32db203

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl requested a review from HumairAK July 14, 2026 17:35

Fix inconsistencies

12f8d72

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl force-pushed the postgres-optimization2 branch from 39fceff to 12f8d72 Compare July 14, 2026 19:29

HumairAK reviewed Jul 14, 2026

View reviewed changes

B-Step62 reviewed Jul 14, 2026

View reviewed changes

Address review feedback

d106a9f

Signed-off-by: mprahl <mprahl@users.noreply.github.com>

mprahl requested review from B-Step62 and HumairAK July 15, 2026 20:42


		### 2. Denormalize hot span cost analytics onto `spans`

		Add the following columns to `spans`:


		Add three daily rollup tables:

		\| Table \| Grain \| Dimension columns \| Measure columns \|

Uh oh!

Conversation

mprahl commented Jul 8, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

etirelli left a comment

Choose a reason for hiding this comment

Uh oh!

mprahl commented Jul 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

etirelli Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

etirelli commented Jul 10, 2026

Uh oh!

mprahl commented Jul 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

HumairAK Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

etirelli Jul 10, 2026 •

edited

Loading

HumairAK Jul 14, 2026 •

edited

Loading