Repository: iterative/dvc

Issue:
dvc exp show --sort-by: fails
# Bug Report

## Description

`dvc exp show --sort-by` throws an error.

### Reproduce

1. dvc init
2. dvc stage add -n foo -M foo.yaml 'echo "score: 1" > foo.yaml'
3. git add .
4. git commit -m "add stage"
5. dvc exp run
6. dvc exp show --no-pager --sort-by score

```console
ERROR: failed to show experiments - Unknown sort column 'score'
```

### Expected

```console
┏━━━━━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━┓
┃ Experiment    ┃ Created  ┃ score ┃
┡━━━━━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━┩
│ workspace     │ -        │     1 │
│ master        │ 04:29 PM │     - │
│ └── exp-985c5 │ 04:29 PM │     1 │
└───────────────┴──────────┴───────┘
```

### Environment information

**Output of `dvc version`:**

```console
$ dvc version
DVC version: 2.0.0a0+adedd1
---------------------------------
Platform: Python 3.9.1 on Linux-5.8.0-38-generic-x86_64-with-glibc2.31
Supports: http, https
Cache types: <https://error.dvc.org/no-dvc-cache>
Caches: local
Remotes: None
Workspace directory: ext4 on /dev/mapper/vgubuntu-root
Repo: dvc, git
```


Discussion:

Resolving patch:
diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -130,14 +130,13 @@ def _collect_rows(
     from dvc.scm.git import Git
 
     if sort_by:
-        if sort_by in metric_names:
-            sort_type = "metrics"
-        elif sort_by in param_names:
-            sort_type = "params"
-        else:
-            raise InvalidArgumentError(f"Unknown sort column '{sort_by}'")
+        sort_path, sort_name, sort_type = _sort_column(
+            sort_by, metric_names, param_names
+        )
         reverse = sort_order == "desc"
-        experiments = _sort_exp(experiments, sort_by, sort_type, reverse)
+        experiments = _sort_exp(
+            experiments, sort_path, sort_name, sort_type, reverse
+        )
 
     new_checkpoint = True
     for i, (rev, exp) in enumerate(experiments.items()):
@@ -193,22 +192,45 @@ def _collect_rows(
         yield row, style
 
 
-def _sort_exp(experiments, sort_by, typ, reverse):
+def _sort_column(sort_by, metric_names, param_names):
+    path, _, sort_name = sort_by.rpartition(":")
+    matches = set()
+
+    if path:
+        if path in metric_names and sort_name in metric_names[path]:
+            matches.add((path, sort_name, "metrics"))
+        if path in param_names and sort_name in param_names[path]:
+            matches.add((path, sort_name, "params"))
+    else:
+        for path in metric_names:
+            if sort_name in metric_names[path]:
+                matches.add((path, sort_name, "metrics"))
+        for path in param_names:
+            if sort_name in param_names[path]:
+                matches.add((path, sort_name, "params"))
+
+    if len(matches) == 1:
+        return matches.pop()
+    if len(matches) > 1:
+        raise InvalidArgumentError(
+            "Ambiguous sort column '{}' matched '{}'".format(
+                sort_by,
+                ", ".join([f"{path}:{name}" for path, name, _ in matches]),
+            )
+        )
+    raise InvalidArgumentError(f"Unknown sort column '{sort_by}'")
+
+
+def _sort_exp(experiments, sort_path, sort_name, typ, reverse):
     def _sort(item):
         rev, exp = item
         tip = exp.get("checkpoint_tip")
         if tip and tip != rev:
             # Sort checkpoint experiments by tip commit
             return _sort((tip, experiments[tip]))
-        for fname, item in exp.get(typ, {}).items():
-            if isinstance(item, dict):
-                item = flatten(item)
-            else:
-                item = {fname: item}
-            if sort_by in item:
-                val = item[sort_by]
-                return (val is None, val)
-        return (True, None)
+        data = exp.get(typ, {}).get(sort_path, {})
+        val = flatten(data).get(sort_name)
+        return (val is None, val)
 
     ret = OrderedDict()
     if "baseline" in experiments:
@@ -472,10 +494,12 @@ def run(self):
 
 class CmdExperimentsRun(CmdRepro):
     def run(self):
+        from dvc.command.metrics import _show_metrics
+
         if self.args.reset:
             logger.info("Any existing checkpoints will be reset and re-run.")
 
-        self.repo.experiments.run(
+        results = self.repo.experiments.run(
             name=self.args.name,
             queue=self.args.queue,
             run_all=self.args.run_all,
@@ -487,6 +511,11 @@ def run(self):
             **self._repro_kwargs,
         )
 
+        if self.args.metrics and results:
+            metrics = self.repo.metrics.show(revs=list(results))
+            metrics.pop("workspace", None)
+            logger.info(_show_metrics(metrics))
+
         return 0
 
 


Test patch:
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -581,3 +581,15 @@ def test_queue(tmp_dir, scm, dvc, exp_stage, mocker):
         with tree.open(tmp_dir / "metrics.yaml") as fobj:
             metrics.add(fobj.read().strip())
     assert expected == metrics
+
+
+def test_run_metrics(tmp_dir, scm, dvc, exp_stage, mocker):
+    from dvc.main import main
+
+    mocker.patch.object(
+        dvc.experiments, "run", return_value={"abc123": "abc123"}
+    )
+    show_mock = mocker.patch.object(dvc.metrics, "show", return_value={})
+
+    main(["exp", "run", "-m"])
+    assert show_mock.called_once()
diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py
--- a/tests/func/experiments/test_show.py
+++ b/tests/func/experiments/test_show.py
@@ -1,3 +1,4 @@
+import logging
 import os
 from datetime import datetime
 
@@ -214,3 +215,21 @@ def test_show_multiple_commits(tmp_dir, scm, dvc, exp_stage):
 
     results = dvc.experiments.show(num=100)
     assert set(results.keys()) == expected
+
+
+def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog):
+    with caplog.at_level(logging.ERROR):
+        assert main(["exp", "show", "--no-pager", "--sort-by=bar"]) != 0
+        assert "Unknown sort column" in caplog.text
+
+    with caplog.at_level(logging.ERROR):
+        assert main(["exp", "show", "--no-pager", "--sort-by=foo"]) != 0
+        assert "Ambiguous sort column" in caplog.text
+
+    assert (
+        main(["exp", "show", "--no-pager", "--sort-by=params.yaml:foo"]) == 0
+    )
+
+    assert (
+        main(["exp", "show", "--no-pager", "--sort-by=metrics.yaml:foo"]) == 0
+    )