RFC: PostgreSQL trace query optimization

[aws-khatria]

Proposes schema-level optimizations to improve MLflow trace analytics performance on PostgreSQL — based on a load test with 10M traces, 30M spans, and 30M assessments in a single experiment.

Foreign key constraints that can be dropped to improve write throughput. The performance improvement load testing is still pending but this can prove to be beneficial for logging high volume metrics in an experiment.

Happy for suggestions to further improve the MLflow dashboarding experience.

[mprahl]

Did you mean something else here? I see an index already for those columns:
https://github.com/mlflow/mlflow/blob/df0bcff70507ee328b9ed73c8781e1d4c21f73a9/mlflow/store/tracking/dbmodels/models.py#L776-L782

[dkuc]

Perhaps they meant covering index? The composite index definitely exists so this is incorrect as written.

[mprahl]

I think the assessments part of this section could be strengthened by avoiding INCLUDE (value) on the covering index and instead introducing precomputed aggregate columns on assessments.

Example schema change:

ALTER TABLE assessments
    ADD COLUMN aggregate_value DOUBLE PRECISION,
    ADD COLUMN aggregate_value_kind SMALLINT;

Where aggregate_value_kind records how aggregate_value was derived, e.g.:

• 1 = numeric
• 2 = boolean
• 3 = yes/no string
• NULL = not aggregateable

Then the numeric/percentile path can operate on aggregate_value directly instead of repeatedly applying CASE / CAST against value. Some of the advantages of doing so:

• faster AVG / Pxx queries because the numeric representation is materialized at write time
• no per-row CASE / CAST work in the hot query path
• avoids carrying value in a covering index since value can contain arbitrary JSON payloads / non-numeric content
• preserves provenance via aggregate_value_kind, so 1.0 from a boolean-like assessment is distinguishable from a true numeric score

Concretely, for the index in this section, perhaps something like this?

CREATE INDEX CONCURRENTLY idx_assessments_traceid_created_ts
    ON assessments (trace_id, created_timestamp)
    INCLUDE (name, aggregate_value, aggregate_value_kind);

And, if this RFC later adopts denormalized experiment_id on assessments, also a partial index for the numeric aggregate path:

CREATE INDEX CONCURRENTLY idx_assessments_exp_ts_name_agg
    ON assessments (experiment_id, last_updated_timestamp, name)
    INCLUDE (aggregate_value, aggregate_value_kind)
    WHERE aggregate_value IS NOT NULL;

I would also add API-side validation when creating a new assessment so we do not mix value types under the same assessment name within the same experiment. Today the aggregate query effectively assumes that a given assessment name has a stable value type/semantics, but that is not enforced. aggregate_value_kind gives us a clean way to validate that invariant and reject writes that would make AVG / Pxx ambiguous or misleading.

What do you think?

[mprahl]

Could we consider storing the trace timestamp (already normalized for the day perhaps) along with the experiment ID so the filter is accurate?

[mprahl]

I think this is also missing cache_read_input_tokens and cache_creation_input_tokens. Is the intention to drop these from the dashboard query or should we include them on the table?

[mprahl]

Can we also add server-side query-shape caps for the trace metrics endpoint?

I ask because max_results only limits the number of grouped rows returned; it does not limit how much data Postgres may scan, join, sort, cast, or aggregate before the LIMIT is applied. Today the UI has a client-side guard that only offers bucket intervals producing at most 1000 chart points, but direct API callers can bypass that and send arbitrary time_interval_seconds values, long time windows, exact percentile aggregations, or query shapes that are much more expensive than the dashboard path. This seems like it could allow unbounded expensive queries against the MLflow / Postgres server even if the RFC makes the common dashboard queries much faster.

Would it make sense to define some server-side request-shape limits as part of this work, for example:

• require time_interval_seconds > 0 and start_time_ms < end_time_ms
• require a bounded time range whenever bucketing is requested
• allow only a fixed set of supported bucket intervals instead of arbitrary values so query cost is predictable
• cap computed bucket count, e.g. ceil((end_time_ms - start_time_ms) / (time_interval_seconds * 1000)) <= 1000
• apply tighter time-window caps for more expensive query shapes
• apply tighter limits or separate handling for exact percentile queries, since percentile_cont can require sorting large groups
• consider constraining very expensive filter/dimension combinations for large windows unless they are explicitly optimized

[mprahl]

I'm supportive of removing this since span insertion is very high volume. The others, I'm not as convinced about. Are you seeing performance issues from those other tables?

[mprahl]

I think just add the ones mentioned here: https://github.com/mlflow/rfcs/pull/23/changes#r3482899422.

[mprahl]

Let's rerun the benchmarks after the optimizations are in place and reconsider it after.

[mprahl]

I say leave out for now.

[mprahl]

Could we also consider denormalizing mlflow.trace.session from the trace_request_metadata table to optimize the dashboard overview's session count?

[mprahl]

I don't think this will allow for the index to be used for average because the backend queries it with a CASE. My other comment had an alternative suggestion for denormalzing the values we need to average and get percentages on.

[mprahl]

From AI:
The proposed index keys on the wrong timestamp column. Section 3 proposes idx_assessments_exp_ts ON assessments (experiment_id, last_updated_timestamp) and rewrites the dashboard filter to last_updated_timestamp. But the assessment time-bucket GROUP BY buckets on created_timestamp, not last_updated_timestamp. See get_time_bucket_expression:

case MetricViewType.ASSESSMENTS:
    timestamp_column = SqlAssessments.created_timestamp

Meanwhile the time-range filter in query_trace_metrics keys off trace_info.timestamp_ms:

if start_time_ms is not None:
    query = query.filter(SqlTraceInfo.timestamp_ms >= start_time_ms)
if end_time_ms is not None:
    query = query.filter(SqlTraceInfo.timestamp_ms <= end_time_ms)

So an index on (experiment_id, last_updated_timestamp) orders neither the created_timestamp bucket/sort nor the timestamp_ms filter — the GROUP BY will still sort. To actually serve the query, key the index on (experiment_id, created_timestamp) to match the real bucketing/ordering. Also worth calling out explicitly in the RFC: the current code already mixes time semantics — it filters on trace_info.timestamp_ms but buckets on assessments.created_timestamp — so the RFC should clarify which timestamp the denormalized index and rewritten filter are meant to align with.

[mprahl]

Would it make sense to store two explicit time axes on assessments: one for when the trace ran (e.g. trace_timestamp_ms, optionally trace_day_bucket_ms) and one for when the assessment was created/updated?

Right now this section is forced to pick between trace-time semantics and assessment-activity semantics. Carrying both would let the API support both questions cleanly: "how did traces that ran in this window score?" and "what scoring activity happened in this window?" It also gives the dashboard a no-join path without baking in ambiguous time semantics.

🤖 Generated with Claude

[mprahl]

I think there may be another worthwhile schema addition specifically for the assessment_count by name+value path.

Even if we materialize numeric aggregate_value for AVG / percentile, this query still groups on raw JSON text in assessments.value, which can be wide and can split logically equivalent payloads by serialization details. Would it help to add a canonicalized grouping representation, e.g. value_text_norm for scalar/enum-like values plus value_hash for larger structured payloads?

Then the count path can group on (name, value_hash) (and optionally display value_text_norm) instead of grouping directly on TEXT JSON.

🤖 Generated with Claude

[mprahl]

Could we widen this section from denormalizing only hot metrics to denormalizing the hot analytics dimensions too?

Even if token columns move onto trace_info, the current trace metrics path still joins trace_tags for trace_name and trace_request_metadata for mlflow.trace.session / session count. A narrow 1:1 analytics projection (either extra columns on trace_info or a sidecar trace_analytics table keyed by request_id) for trace_name, session_id, maybe status, plus the hot token fields, would remove three current joins at once and make future dashboard metrics cheaper to add.

This feels like a better long-term shape than promoting one metric at a time.

🤖 Generated with Claude

[kramaranya]

Thank you, @aws-khatria!
I left a few comments

[kramaranya]

This function doesn't seem to exist, token metrics are written in start_trace() via TOKEN_USAGE metadata. Could you update this please?

[kramaranya]

Nit: RFC also deletes rows from trace_metrics and changes where tokens are written, which isn't purely additive

[kramaranya]

To clarify, do you mean the existing assessment timestamps?

[kramaranya]

What happens if the migration fails after this step?

[dkuc]

One pattern I'm seeing with all the slow time_bucket based queries is that the bucket is being derived dynamically by the group by. If instead we generate the buckets ahead of time with GENERATE_SERIES() and join to it, it should be considerablly faster.

(for these examples I'm using the denormalized table for simplicity)

Before:

SELECT FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000 AS time_bucket,
       SUM(trace_info.input_tokens) AS "SUM"
FROM trace_info
WHERE trace_info.experiment_id = 0
  AND trace_info.timestamp_ms >= 1751916029984
  AND trace_info.timestamp_ms <= 1783452029984
GROUP BY FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000
ORDER BY FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000
LIMIT 1000;

After

SELECT b.bucket_start AS time_bucket,
       sub.sum_val    AS "SUM"
FROM GENERATE_SERIES(1751916029984::bigint, 1783452029984::bigint, 2592000000::bigint) AS b(bucket_start)
         CROSS JOIN LATERAL (
    SELECT SUM(ti.input_tokens) AS sum_val
    FROM trace_info ti
    WHERE ti.experiment_id = 0
      AND ti.timestamp_ms >= b.bucket_start
      AND ti.timestamp_ms < b.bucket_start + 2592000000::bigint
    ) sub
ORDER BY b.bucket_start
LIMIT 1000;

And if we did go with this pattern, the index to go with would probably cover the denormalized columns.

CREATE INDEX idx_trace_info_exp_ts_tokens                                                                                                                                                                                         
      ON trace_info (experiment_id, timestamp_ms)                                                                                                                                                                                   
      INCLUDE (input_tokens, output_tokens, total_tokens);

[dkuc]

Most likely faster to do this without subqueries.

UPDATE trace_info ti                                                                                                                                                                                                              
  SET input_tokens  = tm.input_tokens,                                                                                                                                                                                              
      output_tokens = tm.output_tokens,                                                                                                                                                                                             
      total_tokens  = tm.total_tokens                                                                                                                                                                                               
  FROM (                                                                                                                                                                                                                            
      SELECT request_id,                                                                                                                                                                                                            
             MAX(value) FILTER (WHERE key = 'input_tokens')  AS input_tokens,                                                                                                                                                       
             MAX(value) FILTER (WHERE key = 'output_tokens') AS output_tokens,                                                                                                                                                      
             MAX(value) FILTER (WHERE key = 'total_tokens')  AS total_tokens                                                                                                                                                        
      FROM trace_metrics                                                                                                                                                                                                            
      WHERE key IN ('input_tokens', 'output_tokens', 'total_tokens')                                                                                                                                                                
      GROUP BY request_id                                                                                                                                                                                                           
  ) tm                                                                                                                                                                                                                              
  WHERE ti.request_id = tm.request_id

[mprahl]

One pattern I'm seeing with all the slow time_bucket based queries is that the bucket is being derived dynamically by the group by. If instead we generate the buckets ahead of time with GENERATE_SERIES() and join to it, it should be considerablly faster.

(for these examples I'm using the denormalized table for simplicity)

Before:

SELECT FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000 AS time_bucket,
       SUM(trace_info.input_tokens) AS "SUM"
FROM trace_info
WHERE trace_info.experiment_id = 0
  AND trace_info.timestamp_ms >= 1751916029984
  AND trace_info.timestamp_ms <= 1783452029984
GROUP BY FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000
ORDER BY FLOOR(trace_info.timestamp_ms / CAST(2592000000 AS NUMERIC)) * 2592000000
LIMIT 1000;

After

SELECT b.bucket_start AS time_bucket,
       sub.sum_val    AS "SUM"
FROM GENERATE_SERIES(1751916029984::bigint, 1783452029984::bigint, 2592000000::bigint) AS b(bucket_start)
         CROSS JOIN LATERAL (
    SELECT SUM(ti.input_tokens) AS sum_val
    FROM trace_info ti
    WHERE ti.experiment_id = 0
      AND ti.timestamp_ms >= b.bucket_start
      AND ti.timestamp_ms < b.bucket_start + 2592000000::bigint
    ) sub
ORDER BY b.bucket_start
LIMIT 1000;

And if we did go with this pattern, the index to go with would probably cover the denormalized columns.

CREATE INDEX idx_trace_info_exp_ts_tokens                                                                                                                                                                                         
      ON trace_info (experiment_id, timestamp_ms)                                                                                                                                                                                   
      INCLUDE (input_tokens, output_tokens, total_tokens);

Interesting suggestion!

I tried the GENERATE_SERIES + LATERAL shape on the local Postgres benchmark with ~822k denormalized traces. With the existing (experiment_id, timestamp_ms) index, it was slower than the current grouped query across daily, hourly, and sparse-prefix ranges. The repeated per-bucket index scans and heap lookups were more expensive than Postgres’s parallel scan/group aggregate plan.

I didn’t test the additional covering-index variant proposed here. That may reduce heap lookups, but it would add a large new index and write amplification, so I think I'll skip this for now.

[mprahl]

After a discussion with @aws-khatria offline, we decided that I continue with the RFC due to capacity constraints. The new PR is #24. Thanks so much for your help on this @aws-khatria and thanks for the reviews!