Repository: iterative/dvc

Issue:
Plots: render dots with annotations without Vega spec
If Vega spec is present, Studio render plots with nice annotations with exact value at Y-axis when you hover your mouse over the plot. But if we don't have that spec, Studio won't show them. It would be nice to do that!

## Example with annotations rendered

![Screenshot 2021-08-06 at 14 45 20](https://user-images.githubusercontent.com/6797716/128505857-eb85147a-ac2d-4ff2-b25e-d1f67554ea04.png)

https://studio.iterative.ai/user/aguschin/views/fashion_mnist_project-i4tmmlg2aq

## Example where we do not have annotations

![Screenshot 2021-08-06 at 14 49 09](https://user-images.githubusercontent.com/6797716/128506198-6d255f1f-c1c9-4872-a198-4b668d7a2d35.png)

https://studio.iterative.ai/user/aguschin/views/example-get-started-w3o8p2cq0o

Thanks to @mvshmakov for investigation on the issue origin.


Discussion:
Thanks, @aguschin! @shcheklein is there a priority for this for Studio?
@dberenbaum in Studio plots improvements are considered p1 at the moment. I think it would be great to invest time into improving plots in DVC (as we do with images, live plots already) along the way. Products might affect each other. I would not ruin or disrupt any existing large priorities though for this. In Studio we'll be iterating on basic large features for the next month at least.
Thanks, good to know. Agreed that improving in DVC makes sense. This looks like a very minor change, so it shouldn't disrupt much.
Related comment from @pared about the default template in https://github.com/iterative/dvc/pull/4975/files#r531031121:

> I would consider dropping this one at some point, for sake of interactive linear to become the default, though one thing that is stopping me from proposing it right away is the simplicity of the default plot, which definitely makes the learning curve for custom plots smoother.

I would suggest that we make "linear" the default and rename the current default template to something like "simple" or "basic" and document that it's a simplified linear plot included as an easy starting point for creating custom templates.

The docs only mention the templates in https://dvc.org/doc/command-reference/plots#plot-templates, and the default is listed as "linear plot," with no explanation of the difference between "default.json" and "linear.json," so I think this would clarify and shouldn't create issues for existing users.

Patch:
diff --git a/dvc/render/utils.py b/dvc/render/utils.py
--- a/dvc/render/utils.py
+++ b/dvc/render/utils.py
@@ -1,7 +1,7 @@
 import os.path
 from typing import Dict, List
 
-import dpath
+import dpath.util
 
 
 def get_files(data: Dict) -> List:
diff --git a/dvc/render/vega.py b/dvc/render/vega.py
--- a/dvc/render/vega.py
+++ b/dvc/render/vega.py
@@ -147,7 +147,7 @@ def _datapoints(self, props: Dict):
     def get_vega(self) -> Optional[str]:
         props = self._squash_props()
 
-        template = self.templates.load(props.get("template") or "default")
+        template = self.templates.load(props.get("template", None))
 
         if not props.get("x") and template.has_anchor("x"):
             props["append_index"] = True
diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py
--- a/dvc/repo/plots/__init__.py
+++ b/dvc/repo/plots/__init__.py
@@ -161,7 +161,7 @@ def modify(self, path, props=None, unset=None):
         props = props or {}
         template = props.get("template")
         if template:
-            self.templates.get_template(template)
+            self.templates.load(template)
 
         (out,) = self.repo.find_outs_by_path(path)
         if not out.plot and unset is not None:
diff --git a/dvc/repo/plots/template.py b/dvc/repo/plots/template.py
--- a/dvc/repo/plots/template.py
+++ b/dvc/repo/plots/template.py
@@ -1,15 +1,23 @@
 import json
 import os
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING, Iterable, Optional
+
+try:
+    import importlib_resources
+except ImportError:
+    import importlib.resources as importlib_resources  # type: ignore[no-redef]
 
 from funcy import cached_property
 
 from dvc.exceptions import DvcException
 
+if TYPE_CHECKING:
+    from dvc.types import StrPath
+
 
 class TemplateNotFoundError(DvcException):
-    def __init__(self, path):
-        super().__init__(f"Template '{path}' not found.")
+    def __init__(self, name):
+        super().__init__(f"Template '{name}' not found.")
 
 
 class BadTemplateError(DvcException):
@@ -29,25 +37,9 @@ class Template:
     EXTENSION = ".json"
     ANCHOR = "<DVC_METRIC_{}>"
 
-    DEFAULT_CONTENT: Optional[Dict[str, Any]] = None
-    DEFAULT_NAME: Optional[str] = None
-
-    def __init__(self, content=None, name=None):
-        if content:
-            self.content = content
-        else:
-            self.content = (
-                json.dumps(
-                    self.DEFAULT_CONTENT,
-                    indent=self.INDENT,
-                    separators=self.SEPARATORS,
-                )
-                + "\n"
-            )
-
-        self.name = name or self.DEFAULT_NAME
-        assert self.content and self.name
-        self.filename = self.name + self.EXTENSION
+    def __init__(self, content, name):
+        self.content = content
+        self.name = name
 
     def render(self, data, props=None):
         props = props or {}
@@ -55,7 +47,7 @@ def render(self, data, props=None):
         if self._anchor_str("data") not in self.content:
             anchor = self.anchor("data")
             raise BadTemplateError(
-                f"Template '{self.filename}' is not using '{anchor}' anchor"
+                f"Template '{self.name}' is not using '{anchor}' anchor"
             )
 
         if props.get("x"):
@@ -106,466 +98,86 @@ def _check_field_exists(data, field):
             raise NoFieldInDataError(field)
 
 
-class DefaultTemplate(Template):
-    DEFAULT_NAME = "default"
-
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "width": 300,
-        "height": 300,
-        "mark": {"type": "line"},
-        "encoding": {
-            "x": {
-                "field": Template.anchor("x"),
-                "type": "quantitative",
-                "title": Template.anchor("x_label"),
-            },
-            "y": {
-                "field": Template.anchor("y"),
-                "type": "quantitative",
-                "title": Template.anchor("y_label"),
-                "scale": {"zero": False},
-            },
-            "color": {"field": "rev", "type": "nominal"},
-        },
-    }
-
-
-class ConfusionTemplate(Template):
-    DEFAULT_NAME = "confusion"
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "facet": {"field": "rev", "type": "nominal"},
-        "spec": {
-            "transform": [
-                {
-                    "aggregate": [{"op": "count", "as": "xy_count"}],
-                    "groupby": [Template.anchor("y"), Template.anchor("x")],
-                },
-                {
-                    "impute": "xy_count",
-                    "groupby": ["rev", Template.anchor("y")],
-                    "key": Template.anchor("x"),
-                    "value": 0,
-                },
-                {
-                    "impute": "xy_count",
-                    "groupby": ["rev", Template.anchor("x")],
-                    "key": Template.anchor("y"),
-                    "value": 0,
-                },
-                {
-                    "joinaggregate": [
-                        {"op": "max", "field": "xy_count", "as": "max_count"}
-                    ],
-                    "groupby": [],
-                },
-                {
-                    "calculate": "datum.xy_count / datum.max_count",
-                    "as": "percent_of_max",
-                },
-            ],
-            "encoding": {
-                "x": {
-                    "field": Template.anchor("x"),
-                    "type": "nominal",
-                    "sort": "ascending",
-                    "title": Template.anchor("x_label"),
-                },
-                "y": {
-                    "field": Template.anchor("y"),
-                    "type": "nominal",
-                    "sort": "ascending",
-                    "title": Template.anchor("y_label"),
-                },
-            },
-            "layer": [
-                {
-                    "mark": "rect",
-                    "width": 300,
-                    "height": 300,
-                    "encoding": {
-                        "color": {
-                            "field": "xy_count",
-                            "type": "quantitative",
-                            "title": "",
-                            "scale": {"domainMin": 0, "nice": True},
-                        }
-                    },
-                },
-                {
-                    "mark": "text",
-                    "encoding": {
-                        "text": {"field": "xy_count", "type": "quantitative"},
-                        "color": {
-                            "condition": {
-                                "test": "datum.percent_of_max > 0.5",
-                                "value": "white",
-                            },
-                            "value": "black",
-                        },
-                    },
-                },
-            ],
-        },
-    }
-
-
-class NormalizedConfusionTemplate(Template):
-    DEFAULT_NAME = "confusion_normalized"
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "facet": {"field": "rev", "type": "nominal"},
-        "spec": {
-            "transform": [
-                {
-                    "aggregate": [{"op": "count", "as": "xy_count"}],
-                    "groupby": [Template.anchor("y"), Template.anchor("x")],
-                },
-                {
-                    "impute": "xy_count",
-                    "groupby": ["rev", Template.anchor("y")],
-                    "key": Template.anchor("x"),
-                    "value": 0,
-                },
-                {
-                    "impute": "xy_count",
-                    "groupby": ["rev", Template.anchor("x")],
-                    "key": Template.anchor("y"),
-                    "value": 0,
-                },
-                {
-                    "joinaggregate": [
-                        {"op": "sum", "field": "xy_count", "as": "sum_y"}
-                    ],
-                    "groupby": [Template.anchor("y")],
-                },
-                {
-                    "calculate": "datum.xy_count / datum.sum_y",
-                    "as": "percent_of_y",
-                },
-            ],
-            "encoding": {
-                "x": {
-                    "field": Template.anchor("x"),
-                    "type": "nominal",
-                    "sort": "ascending",
-                    "title": Template.anchor("x_label"),
-                },
-                "y": {
-                    "field": Template.anchor("y"),
-                    "type": "nominal",
-                    "sort": "ascending",
-                    "title": Template.anchor("y_label"),
-                },
-            },
-            "layer": [
-                {
-                    "mark": "rect",
-                    "width": 300,
-                    "height": 300,
-                    "encoding": {
-                        "color": {
-                            "field": "percent_of_y",
-                            "type": "quantitative",
-                            "title": "",
-                            "scale": {"domain": [0, 1]},
-                        }
-                    },
-                },
-                {
-                    "mark": "text",
-                    "encoding": {
-                        "text": {
-                            "field": "percent_of_y",
-                            "type": "quantitative",
-                            "format": ".2f",
-                        },
-                        "color": {
-                            "condition": {
-                                "test": "datum.percent_of_y > 0.5",
-                                "value": "white",
-                            },
-                            "value": "black",
-                        },
-                    },
-                },
-            ],
-        },
-    }
-
-
-class ScatterTemplate(Template):
-    DEFAULT_NAME = "scatter"
-
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "width": 300,
-        "height": 300,
-        "layer": [
-            {
-                "encoding": {
-                    "x": {
-                        "field": Template.anchor("x"),
-                        "type": "quantitative",
-                        "title": Template.anchor("x_label"),
-                    },
-                    "y": {
-                        "field": Template.anchor("y"),
-                        "type": "quantitative",
-                        "title": Template.anchor("y_label"),
-                        "scale": {"zero": False},
-                    },
-                    "color": {"field": "rev", "type": "nominal"},
-                },
-                "layer": [
-                    {"mark": "point"},
-                    {
-                        "selection": {
-                            "label": {
-                                "type": "single",
-                                "nearest": True,
-                                "on": "mouseover",
-                                "encodings": ["x"],
-                                "empty": "none",
-                                "clear": "mouseout",
-                            }
-                        },
-                        "mark": "point",
-                        "encoding": {
-                            "opacity": {
-                                "condition": {
-                                    "selection": "label",
-                                    "value": 1,
-                                },
-                                "value": 0,
-                            }
-                        },
-                    },
-                ],
-            },
-            {
-                "transform": [{"filter": {"selection": "label"}}],
-                "layer": [
-                    {
-                        "encoding": {
-                            "text": {
-                                "type": "quantitative",
-                                "field": Template.anchor("y"),
-                            },
-                            "x": {
-                                "field": Template.anchor("x"),
-                                "type": "quantitative",
-                            },
-                            "y": {
-                                "field": Template.anchor("y"),
-                                "type": "quantitative",
-                            },
-                        },
-                        "layer": [
-                            {
-                                "mark": {
-                                    "type": "text",
-                                    "align": "left",
-                                    "dx": 5,
-                                    "dy": -5,
-                                },
-                                "encoding": {
-                                    "color": {
-                                        "type": "nominal",
-                                        "field": "rev",
-                                    }
-                                },
-                            }
-                        ],
-                    }
-                ],
-            },
-        ],
-    }
-
-
-class SmoothLinearTemplate(Template):
-    DEFAULT_NAME = "smooth"
-
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "mark": {"type": "line"},
-        "encoding": {
-            "x": {
-                "field": Template.anchor("x"),
-                "type": "quantitative",
-                "title": Template.anchor("x_label"),
-            },
-            "y": {
-                "field": Template.anchor("y"),
-                "type": "quantitative",
-                "title": Template.anchor("y_label"),
-                "scale": {"zero": False},
-            },
-            "color": {"field": "rev", "type": "nominal"},
-        },
-        "transform": [
-            {
-                "loess": Template.anchor("y"),
-                "on": Template.anchor("x"),
-                "groupby": ["rev"],
-                "bandwidth": 0.3,
-            }
-        ],
-    }
-
-
-class LinearTemplate(Template):
-    DEFAULT_NAME = "linear"
-
-    DEFAULT_CONTENT = {
-        "$schema": "https://vega.github.io/schema/vega-lite/v4.json",
-        "data": {"values": Template.anchor("data")},
-        "title": Template.anchor("title"),
-        "width": 300,
-        "height": 300,
-        "layer": [
-            {
-                "encoding": {
-                    "x": {
-                        "field": Template.anchor("x"),
-                        "type": "quantitative",
-                        "title": Template.anchor("x_label"),
-                    },
-                    "y": {
-                        "field": Template.anchor("y"),
-                        "type": "quantitative",
-                        "title": Template.anchor("y_label"),
-                        "scale": {"zero": False},
-                    },
-                    "color": {"field": "rev", "type": "nominal"},
-                },
-                "layer": [
-                    {"mark": "line"},
-                    {
-                        "selection": {
-                            "label": {
-                                "type": "single",
-                                "nearest": True,
-                                "on": "mouseover",
-                                "encodings": ["x"],
-                                "empty": "none",
-                                "clear": "mouseout",
-                            }
-                        },
-                        "mark": "point",
-                        "encoding": {
-                            "opacity": {
-                                "condition": {
-                                    "selection": "label",
-                                    "value": 1,
-                                },
-                                "value": 0,
-                            }
-                        },
-                    },
-                ],
-            },
-            {
-                "transform": [{"filter": {"selection": "label"}}],
-                "layer": [
-                    {
-                        "mark": {"type": "rule", "color": "gray"},
-                        "encoding": {
-                            "x": {
-                                "field": Template.anchor("x"),
-                                "type": "quantitative",
-                            }
-                        },
-                    },
-                    {
-                        "encoding": {
-                            "text": {
-                                "type": "quantitative",
-                                "field": Template.anchor("y"),
-                            },
-                            "x": {
-                                "field": Template.anchor("x"),
-                                "type": "quantitative",
-                            },
-                            "y": {
-                                "field": Template.anchor("y"),
-                                "type": "quantitative",
-                            },
-                        },
-                        "layer": [
-                            {
-                                "mark": {
-                                    "type": "text",
-                                    "align": "left",
-                                    "dx": 5,
-                                    "dy": -5,
-                                },
-                                "encoding": {
-                                    "color": {
-                                        "type": "nominal",
-                                        "field": "rev",
-                                    }
-                                },
-                            }
-                        ],
-                    },
-                ],
-            },
-        ],
-    }
-
-
 class PlotTemplates:
     TEMPLATES_DIR = "plots"
-    TEMPLATES = [
-        DefaultTemplate,
-        LinearTemplate,
-        ConfusionTemplate,
-        NormalizedConfusionTemplate,
-        ScatterTemplate,
-        SmoothLinearTemplate,
-    ]
+    PKG_TEMPLATES_DIR = "templates"
 
     @cached_property
     def templates_dir(self):
         return os.path.join(self.dvc_dir, self.TEMPLATES_DIR)
 
-    def get_template(self, path):
-        if os.path.exists(path):
-            return path
-
-        if self.dvc_dir and os.path.exists(self.dvc_dir):
-            t_path = os.path.join(self.templates_dir, path)
-            if os.path.exists(t_path):
-                return t_path
-
-            all_templates = [
+    @staticmethod
+    def _find(templates, template_name):
+        for template in templates:
+            if (
+                template_name == template
+                or template_name + ".json" == template
+            ):
+                return template
+        return None
+
+    def _find_in_project(self, name: str) -> Optional["StrPath"]:
+        if os.path.exists(name):
+            return name
+
+        if os.path.exists(self.templates_dir):
+            templates = [
                 os.path.join(root, file)
                 for root, _, files in os.walk(self.templates_dir)
                 for file in files
             ]
-            matches = [
-                template
-                for template in all_templates
-                if os.path.splitext(template)[0] == t_path
-            ]
-            if matches:
-                assert len(matches) == 1
-                return matches[0]
+            found = self._find(templates, name)
+            if found:
+                return os.path.join(self.templates_dir, found)
+        return None
 
-        raise TemplateNotFoundError(path)
+    @staticmethod
+    def _get_templates() -> Iterable[str]:
+        if (
+            importlib_resources.files(__package__)
+            .joinpath(PlotTemplates.PKG_TEMPLATES_DIR)
+            .is_dir()
+        ):
+            entries = (
+                importlib_resources.files(__package__)
+                .joinpath(PlotTemplates.PKG_TEMPLATES_DIR)
+                .iterdir()
+            )
+            return [entry.name for entry in entries]
+        return []
+
+    @staticmethod
+    def _load_from_pkg(name):
+        templates = PlotTemplates._get_templates()
+        found = PlotTemplates._find(templates, name)
+        if found:
+            return (
+                (
+                    importlib_resources.files(__package__)
+                    / PlotTemplates.PKG_TEMPLATES_DIR
+                    / found
+                )
+                .read_bytes()
+                .decode("utf-8")
+            )
+        return None
+
+    def load(self, name: str = None) -> Template:
+
+        if name is not None:
+            template_path = self._find_in_project(name)
+            if template_path:
+                with open(template_path, "r") as fd:
+                    content = fd.read()
+                return Template(content, name)
+        else:
+            name = "linear"
+
+        content = self._load_from_pkg(name)
+        if content:
+            return Template(content, name)
+
+        raise TemplateNotFoundError(name)
 
     def __init__(self, dvc_dir):
         self.dvc_dir = dvc_dir
@@ -574,24 +186,14 @@ def init(self):
         from dvc.utils.fs import makedirs
 
         makedirs(self.templates_dir, exist_ok=True)
-        for t in self.TEMPLATES:
-            self._dump(t())
-
-    def _dump(self, template):
-        path = os.path.join(self.templates_dir, template.filename)
-        with open(path, "w") as fd:
-            fd.write(template.content)
-
-    def load(self, name):
-        try:
-            path = self.get_template(name)
-
-            with open(path) as fd:
-                content = fd.read()
-
-            return Template(content, name=name)
-        except TemplateNotFoundError:
-            for template in self.TEMPLATES:
-                if template.DEFAULT_NAME == name:
-                    return template()
-            raise
+
+        templates = self._get_templates()
+        for template in templates:
+            content = (
+                importlib_resources.files(__package__)
+                .joinpath(PlotTemplates.PKG_TEMPLATES_DIR)
+                .joinpath(template)
+                .read_text()
+            )
+            with open(os.path.join(self.templates_dir, template), "w") as fd:
+                fd.write(content)


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -161,10 +161,19 @@ def pytest_configure(config):
 
 @pytest.fixture()
 def custom_template(tmp_dir, dvc):
-    import shutil
+    try:
+        import importlib_resources
+    except ImportError:
+        import importlib.resources as importlib_resources
+
+    content = (
+        importlib_resources.files("dvc.repo.plots")
+        / "templates"
+        / "simple.json"
+    ).read_text()
 
     template = tmp_dir / "custom_template.json"
-    shutil.copy(tmp_dir / ".dvc" / "plots" / "default.json", template)
+    template.write_text(content)
     return template
 
 
diff --git a/tests/unit/render/test_render.py b/tests/unit/render/test_render.py
--- a/tests/unit/render/test_render.py
+++ b/tests/unit/render/test_render.py
@@ -63,5 +63,9 @@ def test_render(tmp_dir, dvc):
     index_content = index_path.read_text()
     file_vega = find_vega(dvc, data, "file.json")
     some_vega = find_vega(dvc, data, "some.csv")
-    assert file_vega in index_content.strip()
-    assert some_vega in index_content.strip()
+
+    def clean(txt: str) -> str:
+        return txt.replace("\n", "").replace("\r", "").replace(" ", "")
+
+    assert clean(file_vega) in clean(index_content)
+    assert clean(some_vega) in clean(index_content)
diff --git a/tests/unit/render/test_vega.py b/tests/unit/render/test_vega.py
--- a/tests/unit/render/test_vega.py
+++ b/tests/unit/render/test_vega.py
@@ -3,6 +3,7 @@
 from collections import OrderedDict
 
 import pytest
+from funcy import first
 
 from dvc.render.utils import find_vega, group_by_filename
 from dvc.render.vega import (
@@ -117,10 +118,12 @@ def test_one_column(tmp_dir, scm, dvc):
         {"val": 2, INDEX_FIELD: 0, REVISION_FIELD: "workspace"},
         {"val": 3, INDEX_FIELD: 1, REVISION_FIELD: "workspace"},
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "val"
-    assert plot_content["encoding"]["x"]["title"] == "x_title"
-    assert plot_content["encoding"]["y"]["title"] == "y_title"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "val"
+    assert first(plot_content["layer"])["encoding"]["x"]["title"] == "x_title"
+    assert first(plot_content["layer"])["encoding"]["y"]["title"] == "y_title"
 
 
 def test_multiple_columns(tmp_dir, scm, dvc):
@@ -152,8 +155,10 @@ def test_multiple_columns(tmp_dir, scm, dvc):
             "second_val": 300,
         },
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "val"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "val"
 
 
 def test_choose_axes(tmp_dir, scm, dvc):
@@ -184,8 +189,12 @@ def test_choose_axes(tmp_dir, scm, dvc):
             "second_val": 300,
         },
     ]
-    assert plot_content["encoding"]["x"]["field"] == "first_val"
-    assert plot_content["encoding"]["y"]["field"] == "second_val"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == "first_val"
+    )
+    assert (
+        first(plot_content["layer"])["encoding"]["y"]["field"] == "second_val"
+    )
 
 
 def test_confusion(tmp_dir, dvc):
@@ -250,8 +259,10 @@ def test_multiple_revs_default(tmp_dir, scm, dvc):
         {"y": 2, INDEX_FIELD: 0, REVISION_FIELD: "v1"},
         {"y": 3, INDEX_FIELD: 1, REVISION_FIELD: "v1"},
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "y"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "y"
 
 
 def test_metric_missing(tmp_dir, scm, dvc, caplog):
@@ -270,8 +281,10 @@ def test_metric_missing(tmp_dir, scm, dvc, caplog):
         {"y": 2, INDEX_FIELD: 0, REVISION_FIELD: "v2"},
         {"y": 3, INDEX_FIELD: 1, REVISION_FIELD: "v2"},
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "y"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "y"
 
 
 def test_custom_template(tmp_dir, scm, dvc, custom_template):
@@ -353,8 +366,10 @@ def test_plot_default_choose_column(tmp_dir, scm, dvc):
         {INDEX_FIELD: 0, "b": 2, REVISION_FIELD: "workspace"},
         {INDEX_FIELD: 1, "b": 3, REVISION_FIELD: "workspace"},
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "b"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "b"
 
 
 def test_raise_on_wrong_field(tmp_dir, scm, dvc):
@@ -432,5 +447,7 @@ def test_find_vega(tmp_dir, dvc):
         {"y": 2, INDEX_FIELD: 0, REVISION_FIELD: "v1"},
         {"y": 3, INDEX_FIELD: 1, REVISION_FIELD: "v1"},
     ]
-    assert plot_content["encoding"]["x"]["field"] == INDEX_FIELD
-    assert plot_content["encoding"]["y"]["field"] == "y"
+    assert (
+        first(plot_content["layer"])["encoding"]["x"]["field"] == INDEX_FIELD
+    )
+    assert first(plot_content["layer"])["encoding"]["y"]["field"] == "y"