feat(tracking): Add composite index on metrics (run_uuid, key, step) (mlflow#21727)

Signed-off-by: Raviprakash Darbha <rdarbha@amazon.com>
Signed-off-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com>
Co-authored-by: Raviprakash Darbha <rdarbha@amazon.com>
Co-authored-by: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com>

Author: 3 people

mlflow/store/db_migrations/versions/a5b4c3d2e1f0_add_metrics_run_key_step_index.py

@@ -0,0 +1,23 @@
+"""add composite index on metrics (run_uuid, key, step)
+
+Create Date: 2026-03-10 20:46:00.000000
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a5b4c3d2e1f0"
+down_revision = "c3d6457b6d8a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Add composite index to speed up metric history queries that filter by
+    # run_uuid and key, and order by step. See https://github.com/mlflow/mlflow/issues/12813
+    op.create_index("index_metrics_run_uuid_key_step", "metrics", ["run_uuid", "key", "step"])
+
+
+def downgrade():
+    op.drop_index("index_metrics_run_uuid_key_step", table_name="metrics")

mlflow/store/tracking/dbmodels/models.py

@@ -398,6 +398,7 @@ class SqlMetric(Base):
             "key", "timestamp", "step", "run_uuid", "value", "is_nan", name="metric_pk"
         ),
         Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+        Index(f"index_{__tablename__}_run_uuid_key_step", "run_uuid", "key", "step"),
     )
 
     key = Column(String(250))

tests/store/tracking/test_sqlalchemy_store_schema.py

@@ -147,6 +147,38 @@ def test_create_index_on_run_uuid(tmp_path, db_url):
         assert run_uuid_index_names.issubset(all_index_names)
 
 
+def test_create_index_on_metrics_run_uuid_key_step_via_migration(tmp_path, db_url):
+    # Test that the index is created when upgrading from initial schema via migrations
+    engine = sqlalchemy.create_engine(db_url)
+    InitialBase.metadata.create_all(engine)
+    invoke_cli_runner(mlflow.db.commands, ["upgrade", db_url])
+    with sqlite3.connect(db_url[len("sqlite:///") :]) as conn:
+        cursor = conn.cursor()
+        cursor.execute("SELECT name FROM sqlite_master WHERE type = 'index'")
+        all_index_names = [r[0] for r in cursor.fetchall()]
+        assert "index_metrics_run_uuid_key_step" in all_index_names
+
+        cursor.execute("PRAGMA index_info('index_metrics_run_uuid_key_step')")
+        columns = [row[2] for row in cursor.fetchall()]
+        assert columns == ["run_uuid", "key", "step"]
+
+
+def test_create_index_on_metrics_run_uuid_key_step(tmp_path, db_url):
+    # Test for
+    # mlflow/store/db_migrations/versions/a5b4c3d2e1f0_add_metrics_run_key_step_index.py
+    SqlAlchemyStore(db_url, tmp_path.joinpath("ARTIFACTS").as_uri())
+    with sqlite3.connect(db_url[len("sqlite:///") :]) as conn:
+        cursor = conn.cursor()
+        cursor.execute("SELECT name FROM sqlite_master WHERE type = 'index'")
+        all_index_names = [r[0] for r in cursor.fetchall()]
+        assert "index_metrics_run_uuid_key_step" in all_index_names
+
+        # Verify the index columns and order
+        cursor.execute("PRAGMA index_info('index_metrics_run_uuid_key_step')")
+        columns = [row[2] for row in cursor.fetchall()]
+        assert columns == ["run_uuid", "key", "step"]
+
+
 def test_index_for_dataset_tables(tmp_path, db_url):
     # Test for
     # mlflow/store/db_migrations/versions/7f2a7d5fae7d_add_datasets_inputs_input_tags_tables.py