From d8068430e7c821ca374eba2fd48c992c000f23b9 Mon Sep 17 00:00:00 2001
From: dberenbaum <dave@iterative.ai>
Date: Mon, 28 Nov 2022 12:48:57 -0500
Subject: [PATCH] plots: fix multi-file plots

---
 dvc/render/converter/vega.py             | 29 +++++++++++-------
 tests/unit/render/test_vega_converter.py | 38 ++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/dvc/render/converter/vega.py b/dvc/render/converter/vega.py
index e9c302f7f..55121b163 100644
--- a/dvc/render/converter/vega.py
+++ b/dvc/render/converter/vega.py
@@ -59,13 +59,12 @@ def _verify_field(file2datapoints: Dict[str, List], filename: str, field: str):
     return
 
 
-def _get_x(properties: Dict, file2datapoints: Dict[str, List[Dict]]):
+def _get_xs(properties: Dict, file2datapoints: Dict[str, List[Dict]]):
     x = properties.get("x", None)
     if x is not None and isinstance(x, dict):
-        filename, field = first(_file_field(x))
-        _verify_field(file2datapoints, filename, field)
-        return filename, field
-    return None
+        for filename, field in _file_field(x):
+            _verify_field(file2datapoints, filename, field)
+            yield filename, field
 
 
 def _get_ys(properties, file2datapoints: Dict[str, List[Dict]]):
@@ -193,20 +192,28 @@ class VegaConverter(Converter):
 
         props_update = {}
 
-        x = _get_x(properties, file2datapoints)
+        xs = list(_get_xs(properties, file2datapoints))
 
         # assign "step" if no x provided
-        if not x:
+        if not xs:
             x_file, x_field = (
                 None,
                 INDEX_FIELD,
             )
         else:
-            x_file, x_field = x
+            x_file, x_field = xs[0]
         props_update["x"] = x_field
 
         ys = list(_get_ys(properties, file2datapoints))
 
+        num_xs = len(xs)
+        num_ys = len(ys)
+        if num_xs > 1 and num_xs != num_ys:
+            raise DvcException(
+                f"Cannot have different number of x and y data sources. Found "
+                f"{num_xs} x and {num_ys} y data sources."
+            )
+
         all_datapoints = []
         all_y_fields = {y_field for _, y_field in ys}
 
@@ -216,7 +223,9 @@ class VegaConverter(Converter):
         else:
             props_update["y"] = first(all_y_fields)
 
-        for y_file, y_field in ys:
+        for i, (y_file, y_field) in enumerate(ys):
+            if num_xs > 1:
+                x_file, x_field = xs[i]
             datapoints = deepcopy(file2datapoints.get(y_file, []))
 
             if props_update.get("y", None) == "dvc_inferred_y_value":
@@ -239,7 +248,7 @@ class VegaConverter(Converter):
                 except IndexError:
                     raise DvcException(
                         f"Cannot join '{x_field}' from '{x_file}' and "
-                        "'{y_field}' from '{y_file}'. "
+                        f"'{y_field}' from '{y_file}'. "
                         "They have to have same length."
                     )
 
diff --git a/tests/unit/render/test_vega_converter.py b/tests/unit/render/test_vega_converter.py
index 64179279d..a68a1156d 100644
--- a/tests/unit/render/test_vega_converter.py
+++ b/tests/unit/render/test_vega_converter.py
@@ -235,6 +235,44 @@ def test_finding_lists(dictionary, expected_result):
             },
             id="multi_source_json",
         ),
+        pytest.param(
+            {
+                "f": {"metric": [{"v": 1, "v2": 0.1}, {"v": 2, "v2": 0.2}]},
+                "f2": {"metric": [{"v": 3, "v2": 0.3}]},
+            },
+            {"x": "v", "y": {"f": "v2", "f2": "v2"}},
+            [
+                {
+                    "v": 1,
+                    "v2": 0.1,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v2",
+                    },
+                },
+                {
+                    "v": 2,
+                    "v2": 0.2,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v2",
+                    },
+                },
+                {
+                    "v": 3,
+                    "v2": 0.3,
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f2",
+                        "field": "v2",
+                    },
+                },
+            ],
+            {"x": "v", "y": "v2", "x_label": "v", "y_label": "v2"},
+            id="multi_file_json",
+        ),
     ],
 )
 def test_convert(